如何创建可以包装协程或函数的Python装饰器？ [英] How to create a Python decorator that can wrap either coroutine or function?

问题描述

我试图制作一个装饰器来包装协程或函数。

I am trying to make a decorator to wrap either coroutines or functions.

我尝试的第一件事是在包装器中编写一个简单的重复代码：

The first thing I tried was a simple duplicate code in wrappers:

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

这行得通，但我想避免重复代码，因为这并不比编写两个单独的装饰器好得多。

This works, but i want to avoid duplication of code, as this is not much better than writing two separate decorators.

然后我尝试使用类制作装饰器：

Then i tried to make a decorator using class:

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

在某些情况下对我来说效果很好，但是在此解决方案中，我无法在带有或 try 语句中放置函数。
有什么方法可以创建装饰器而不重复代码吗？

That works in some cases very well for me, but in this solution i can't put a function in with or try statements. Is there any way i can create a decorator without duplicating code?

推荐答案

也许您可以找到更好的方法做到这一点，但是例如，您可以将包装逻辑移至某个上下文管理器以防止代码重复：

May be you can find better way to do it, but, for example, you can just move your wrapping logic to some context manager to prevent code duplication:

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

这篇关于如何创建可以包装协程或函数的Python装饰器？的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！