带有可选参数的Python装饰器(这是函数)

时间:2014-11-27 09:38:56

标签: python python-2.7 decorator python-decorators

注意:我知道带有可选参数的装饰器包含三个嵌套函数。但这里的可选参数是功能本身。在将此标记为重复之前,请仔细阅读完整的帖子。我已经尝试了带有可选参数的装饰器的所有技巧,但我找不到任何以函数作为参数。

我有一个包装错误的装饰器:

def wrap_error(func):
    from functools import wraps

    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except:
            import sys

            exc_msg = traceback.format_exception(*sys.exc_info())
            raise MyCustomError(exc_msg)

    return wrapper

如果某个函数引发任何异常,它会包装错误。这个包装使用如下:

@wrap_error
def foo():
    ...

现在我想用附加的回调函数修改这个包装器,这个函数是可选的。我希望这个包装器用作:

@wrap_error
def foo():
    ...

@wrap_error(callback)
def foo():
    ...

我知道如何使用可选参数编​​写装饰器(如果传递的参数不是函数,则基于包装器中的isfunction(func)检查)。但我不知道如何处理这种情况。

注意:我可以使用@wrap_error()代替@wrap_error。此包装器用于多个包中,并且无法更新所有

中的更改

以下是拦截器: 将包装器视为:

@wrap_error(callback)               --->       foo = wrap_error(callback)(foo)
def foo():
    ...

因此,在执行wrap_error(foo)时,我们不知道是否会有任何回调函数用于执行(如果我们仅使用@wrap_error而不是@wrap_error(callback) })。

如果没有(callback),wrap_error中的包装函数将返回func(*args. **kwargs),以便我可以引发异常。否则,我们必须返回func以便在下一步调用它,如果func()引发异常,我们会在callback()块中调用except

2 个答案:

答案 0 :(得分:2)

要在尝试回答问题之前总结问题,您需要一个在以下两个上下文中都能正常工作的装饰器:

@decorator  # case 1
def some_func(...):
    ...

@decorator(some_callback)  # case 2
def some_func(...):
    ...

或者,要展开@语法以澄清事情:

some_func = decorator(some_func)  # case 1

some_func = decorator(some_callback)(some_func)  # case 2

正如我所看到的,这里的棘手问题是decorator很难区分some_funcsome_callback之间的区别(因此在案例1之间)和2);两者都是(可能)只是可调用的对象。


一个可能的解决方案是提供命名参数:

# imports at top of file, not in function definitions
from functools import wraps
import sys

def decorator(func=None, callback=None):
    # Case 1
    if func is not None:
        @wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)  # or whatever
        return wrapper
    # Case 2
    elif callback is not None: 
        def deco(f):
            @wraps(f)
            def wrapper(*args, **kwargs):
                return callback(f(*args, **kwargs))  # or whatever
            return wrapper
        return deco

这使得案例2看起来略有不同:

@decorator(callback=some_callback)
def some_func(...):
    ...

但除此之外你想做什么。请注意,您说您不能使用的选项

@decorator()
def some_func(...):
    ...

不能使用此功能,因为装饰器需要提供funccallback(否则将返回None,这是不可调用的,所以你得到TypeError)。

答案 1 :(得分:2)

由于很难从decorator(func)告诉decorator(callback),请制作两个装饰器:

from functools import wraps

class MyCustomError(Exception):
    def __init__(self):
        print('in MyCustomError')

# Common implementation
def wrap(func,cb=None):
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except:
            if cb is not None:
                cb()
            raise MyCustomError()
    return wrapper

# No parameters version
def wrap_error(func):
    return wrap(func)

# callback parameter version
def wrap_error_cb(cb):
    def deco(func):
        return wrap(func,cb)
    return deco

@wrap_error
def foo(a,b):
    print('in foo',a,b)
    raise Exception('foo exception')

def callback():
    print('in callback')

@wrap_error_cb(callback)
def bar(a):
    print('in bar',a)
    raise Exception('bar exception')

检查foo和bar是否正确使用functools.wraps

>>> foo
<function foo at 0x0000000003F00400>
>>> bar
<function bar at 0x0000000003F00598>

检查包装的函数是否有效:

>>> foo(1,2)
in foo 1 2
in MyCustomError
Traceback (most recent call last):
  File "<interactive input>", line 1, in <module>
  File "C:\test.py", line 16, in wrapper
    raise MyCustomError()
MyCustomError
>>> bar(3)
in bar 3
in callback
in MyCustomError
Traceback (most recent call last):
  File "<interactive input>", line 1, in <module>
  File "C:\test.py", line 16, in wrapper
    raise MyCustomError()
MyCustomError

更新

这是使用您请求的语法来实现的方法,但我认为上述答案更清晰。

from functools import wraps

class MyCustomError(Exception):
    def __init__(self):
        print('in MyCustomError')

# Common implementation
def wrap(func,cb=None):
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except:
            if cb is not None:
                cb()
            raise MyCustomError()
    return wrapper

def wrap_error(func_or_cb):
    # If the function is tagged as a wrap_error_callback
    # return a decorator that returns the wrapped function
    # with a callback.
    if hasattr(func_or_cb,'cb'):
        def deco(func):
            return wrap(func,func_or_cb)
        return deco
    # Otherwise, return a wrapped function without a callback.
    return wrap(func_or_cb)

# decorator to tag callbacks so wrap_error can distinguish them
# from *regular* functions.
def wrap_error_callback(func):
    func.cb = True
    return func

### Examples of use

@wrap_error
def foo(a,b):
    print('in foo',a,b)
    raise Exception('foo exception')

@wrap_error_callback
def callback():
    print('in callback')

@wrap_error(callback)
def bar(a):
    print('in bar',a)
    raise Exception('bar exception')