如何創建一個可以包裝協程或函數的Python裝飾器？

Question

我想讓一個裝飾器來包裝協程或函數。

我想的第一件事是在包裝的簡單重複的代碼：

def duration(func): 
    @functools.wraps(func) 
    def wrapper(*args, **kwargs): 
     start_ts = time.time() 
     result = func(*args, **kwargs) 
     dur = time.time() - start_ts 
     print('{} took {:.2} seconds'.format(func.__name__, dur)) 
     return result 

    @functools.wraps(func) 
    async def async_wrapper(*args, **kwargs): 
     start_ts = time.time() 
     result = await func(*args, **kwargs) 
     dur = time.time() - start_ts 
     print('{} took {:.2} seconds'.format(func.__name__, dur)) 
     return result 

    if asyncio.iscoroutinefunction(func): 
     return async_wrapper 
    else: 
     return wrapper 

這工作，但我想避免重複代碼，因爲這不是比寫兩個單獨的裝飾好得多。

然後我試圖做一個裝飾用類：

class SyncAsyncDuration: 
    def __init__(self): 
     self.start_ts = None 

    def __call__(self, func): 
     @functools.wraps(func) 
     def sync_wrapper(*args, **kwargs): 
      self.setup(func, args, kwargs) 
      result = func(*args, **kwargs) 
      self.teardown(func, args, kwargs) 
      return result 

     @functools.wraps(func) 
     async def async_wrapper(*args, **kwargs): 
      self.setup(func, args, kwargs) 
      result = await func(*args, **kwargs) 
      self.teardown(func, args, kwargs) 
      return result 

     if asyncio.iscoroutinefunction(func): 
      return async_wrapper 
     else: 
      return sync_wrapper 

    def setup(self, func, args, kwargs): 
     self.start_ts = time.time() 

    def teardown(self, func, args, kwargs): 
     dur = time.time() - self.start_ts 
     print('{} took {:.2} seconds'.format(func.__name__, dur)) 

，在某些情況下，工作對我非常好，但在這個解決方案，我不能把一個函數與或嘗試聲明。 有沒有什麼辦法可以創建一個裝飾器而不需要複製代碼？

Answer 1

可能是你能找到更好的方式來做到這一點，但是，例如，你可以將你的包裝邏輯一些上下文管理，防止重複代碼：

import asyncio 
import functools 
import time 
from contextlib import contextmanager 


def duration(func): 
    @contextmanager 
    def wrapping_logic(): 
     start_ts = time.time() 
     yield 
     dur = time.time() - start_ts 
     print('{} took {:.2} seconds'.format(func.__name__, dur)) 

    @functools.wraps(func) 
    def wrapper(*args, **kwargs): 
     if not asyncio.iscoroutinefunction(func): 
      with wrapping_logic(): 
       return func(*args, **kwargs) 
     else: 
      async def tmp(): 
       with wrapping_logic(): 
        return (await func(*args, **kwargs)) 
      return tmp() 
    return wrapper