Comment créer un décorateur Python qui peut envelopper la coroutine ou la fonction?

Question

J'essaye de faire un décorateur pour envelopper des coroutines ou des fonctions.

La première chose que j'ai essayé était un simple code en double dans des emballages:

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 

Cela fonctionne, mais je veux éviter la duplication de code, car ce n'est pas beaucoup mieux que d'écrire deux décorateurs séparés.

Alors j'ai essayé de faire un décorateur classe en utilisant:

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)) 

qui fonctionne dans certains cas, très bien pour moi, mais dans cette solution, je ne peux pas mettre une fonction avec ou essayer déclarations. Puis-je créer un décorateur sans dupliquer le code?

Answer 1

Peut être que vous pouvez trouver une meilleure façon de le faire, mais, par exemple, vous pouvez simplement déplacer votre logique d'emballage à un certain gestionnaire de contexte pour éviter la duplication de code:

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