scribble

ottocho's blog

Home About GitHub

05 Aug 2014
Decorator With Optional Arguments

带可选参数的装饰器

如下一个装饰器实现,使得函数返回值再加上 inc 的值:

import functools

def plus1(method, inc=5):
    @functools.wraps(method)
    def _wrap(*args, **kwargs):
        return method(*args, **kwargs) + inc
    return _wrap

如下示例代码:

@plus1
def k(i):
    return i
print k(5)

# output
# 10

正常输出。但是在如下情况:

@plus1(inc=10)
def k(i):
    return i

print k(5)

会输出如下错误:

Traceback (most recent call last):
  File "yp.py", line 1111, in <module>
    @plus1(inc=10)
TypeError: plus1() takes at least 1 argument (1 given)

这是由于在 @plus(inc=10) 装饰函数时缺失了 method

因此此时自然而然想到再包一层返回函数:

def plus2(inc=5):
    def _plus(method):
        @functools.wraps(method)
        def _wrap(*args, **kwargs):
            return method(*args, **kwargs) + inc
        return _wrap
    return _plus

此时:

@plus2(inc=10)
def k(i):
    return i

print k(5)
# output
# 15

输出正常,但是其实此时 plus2 它不是装饰器,而是 返回装饰器的函数。因此,在不带参数的情况下:

@plus2
def k(i):
    return i

这样的使用是错误的,提示:

  File "/usr/lib/python2.7/functools.py", line 33, in update_wrapper
    setattr(wrapper, attr, getattr(wrapped, attr))
AttributeError: 'int' object has no attribute '__module__'

正确使用应该:

@plus2()
def k(i):
    return i

print k(5)
# output
# 10

plus2() 返回了一个装饰器,因此可以被用以装饰函数。

这样真的很丑。因此我们使用一下 partial,来完成这个事情:

import functools

def plus3(method=None, inc=5):
    if method is None:
        return functools.partial(plus3, inc=inc)
    @functools.wraps(method)
    def _wrap(*args, **kwargs):
        return method(*args, **kwargs) + inc
    return _wrap

@plus3
def k(i):
    return i

@plus3(inc=10)
def j(i):
    return i

print k(5)
# output
# 10

print j(5)
# output
# 15

利用 partial 直接构造一个新的装饰器函数,使其直接兼容参数的两种情况。


Til next time,
at 22:17

scribble

Home About GitHub