如何创建一个可以包装协程或函数的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