Python decorator & wrapt 库

前言

Python 的装饰器是个好东西，自从加入了语法糖 @ 之后。
装饰器方法可以实现很多功能，同时又不会把代码给弄乱。
这里会剖析何为装饰器，并且推荐两个简化装饰器的库

什么是装饰器

装饰器可以在函数裹上一层壳，在函数调度前和调度后触发新的逻辑。
这样可以在原有的逻辑不动的情况下，添加新的逻辑。
比如修改函数的传参，修改函数的返回值。
也可以通过装饰器给原来简单的函数添加新的功能。


def deco_func(func):
    def decorator(*args, **kwargs):
        print("before call")
        res = func(*args, **kwargs)
        print("after call")
        return res

    return decorator


@deco_func
def call_func():
    print("on running")


call_func()
# 打印如下
# before call
# on running
# after call

装饰器的本质其实下面这种写法的语法糖。


def call_func():
    print("on running")

call_func = deco_func(call_func)

两者的本质是一样的，但是用语法糖可以更加清晰看到这个函数被什么所包裹。
另外实现装饰器的方法还可以用类来实现。

class DecoratorClass(object):
    def __init__(self,func):
        self.func = func
    
    def __call__(self,*args, **kwargs):
        print("before call")
        result=  self.func(*args, **kwargs)
        print("after call")
        return result

上面实现的装饰器不能携带参数，如果需要传入参数需要返回一个装饰器

def deco_func(arg1=1,arg2=2):
    def decorator(func):
        def wrapper(*args, **kwargs):
            print(arg1,arg2)
            print("before call")
            res = func(*args, **kwargs)
            print("after call")
            return res

        return wrapper
    return decorator

@deco_func(arg1=3)
def call_func():
    print("on running")


call_func()
# 3 2
# before call
# on running
# after call

这里需要三重函数嵌套会让装饰器繁琐且复杂。
而且由于装饰器返回的函数和源函数是不一致的，默认情况下无法继承 __doc__ 等函数属性。
解决这个问题需要给嵌套里面的函数加上 `@functools.wraps(func)` 的装饰器来同步才行。

import functools

def deco_func(func):
  @functools.wraps(func)
    def decorator(*args, **kwargs):
        return func(*args, **kwargs)
    return decorator

decorator 库介绍

想要解决这个问题有两个库可以推荐 decorator & wrapt

Python 库	wrapt	decorator
类装饰	✔	❌
受欢迎程度	✔	❌
简洁	❌	✔
性能	❌	✔

wrapt 解决的问题更多，比 decorator 更加流行。
decorator 则使用简单且方便，而且性能比 wrapt 要好。
可以根据实际的使用情况来选择。

装饰案例

import timeit
from decorator import decorator

@decorator
def warn_slow(func, timelimit=60, *args, **kw):
    t0 = time.time()
    result = func(*args, **kw)
    dt = time.time() - t0
    if dt > timelimit:
        logging.warn('%s took %d seconds', func.__name__, dt)
    else:
        logging.info('%s took %d seconds', func.__name__, dt)
    return result

@warn_slow  # warn if it takes more than 1 minute
def preprocess_input_files(inputdir, tempdir):
    ...

@warn_slow(timelimit=600)  # warn if it takes more than 10 minutes
def run_calculation(tempdir, outdir):
    ...

from decorator import decorator
@decorator
def with_arguments(func, myarg1=2, myarg2=3, *args, **kwargs):
    return func(*args, **kwargs)

@with_arguments(myarg1=1,myarg2=2)
def function():
    pass

decorator 库只需要一个装饰器就可以让函数同时支持带参数和不带参数。

import wrapt

def with_arguments(myarg1, myarg2):
    @wrapt.decorator
    def wrapper(wrapped, instance, args, kwargs):
        return wrapped(*args, **kwargs)
    return wrapper

@with_arguments(1, 2)
def function():
    pass

wrapt 库则还是需要使用嵌套函数才能实现，但是可以简化一个函数嵌套。

注意事项

from decorator import decorator

@decorator
def pass_through(func, timelimit=60, *args, **kw):
    return func(*args, **kw)

@pass_through
class Class(object):
    @staticmethod
    def function_im(cls, *args):
        print(cls)

# Traceback (most recent call last):
#   File "<stdin>", line 2, in <module>
#   File "<decorator-gen-1>", line 3, in pass_through
#   File "c:\_thm\rez_local_cache\ext\decorator\4.4.2\site-packages\decorator.py", line 252, in decorate        
#     evaldict, __wrapped__=func)
#   File "c:\_thm\rez_local_cache\ext\decorator\4.4.2\site-packages\decorator.py", line 213, in create
#     self = cls(func, name, signature, defaults, doc, module)
#   File "c:\_thm\rez_local_cache\ext\decorator\4.4.2\site-packages\decorator.py", line 143, in __init__        
#     raise TypeError('You are decorating a non function: %s' % func)
# TypeError: You are decorating a non function: <class '__main__.Class'>

decorator 不能装饰类，但是 wrapt 可以做到

import wrapt

@wrapt.decorator
def pass_through(wrapped, instance, args, kwargs):
    print(wrapped,instance)
    res = wrapped(*args, **kwargs)
    return res

class Base(object):
    def __init__(self):
        print("call Base")

@pass_through
class Class(Base):
    def __init__(self):
        super(Class.__wrapped__, self).__init__()

    def function_im(self, *args):
        print(self)

print(type(Class))  # <class 'FunctionWrapper'>
a = Class()
a.function_im(1)