Python 装饰器

装饰器

先来看一个代码段：

def announce(f):
    def wrapper(name):
        print("start to run the function...")
        f(name)
        print("Done with the function.")
    return wrapper

@announce
def hello(name):
    print(f"hello, {name}")

name = "Sudley"
print(f"function name: {hello.__name__}")
hello(name)

@announce就是装饰器，通过装饰器可以快速的给函数添加功能。装饰器让你在一个函数的前后去执行代码。
上面是装饰器可以理解为hello=announce(hello)，所以hello(name)相当于wrapper(name)
上面查看hello的函数名称是announce中返回的函数名称wrapper，也行你想hello依然是hello可以使用functools.wraps来修饰一下，@wraps接受一个函数来进行装饰，并加入了复制函数名称、注释文档、参数列表等等的功能。这可以让我们在装饰器里面访问在装饰之前的函数的属性。

from functools import wraps

def announce(f):
    @wraps(f)
    def wrapper(name):
        print("start to run the function...")
        f(name)
        print("Done with the function.")
    return wrapper

@announce
def hello(name):
    print(f"hello, {name}")

name = "Sudley"
print(f"function name: {hello.__name__}")
hello(name)

装饰器常见用法举例

1.日志（logging）：

平常可能会遇到需要定义函数开始运行和介绍运行的日志，如

import logging
logging.basicConfig(level=logging.INFO)
def square_func(x):
    logging.info("square_func was called")
    return x * x

def add_func(x, y):
    logging.info("add_func was called")
    return x + y

s_result = square_func(4)
a_result = add_func(2, 3)
print(s_result, a_result)

可以看到上面的logging.info有一点坏味道，代码重复了，这时就可以用到装饰器进行优化，如下

from functools import wraps
import logging
logging.basicConfig(level=logging.INFO)
def logit(func):
    @wraps(func)
    def with_called_logging(*args, **kwargs):
        logging.info(func.__name__ + " was called")
        return func(*args, **kwargs)
    return with_called_logging

@logit
def square_func(x):
    return x * x

@logit
def add_func(x, y):
    return x + y

s_result = square_func(4)
a_result = add_func(2, 3)
print(s_result, a_result)

2.函数或方法失败重试（retry）：

当通过wget从网站上download文件时，也许会因为服务、网络不稳定或者其他原因导致失败，也许是遇到性能瓶颈或者暂时无法找到根因，这时sleep一段时间并retry也许能规避当前问题

import time

def retry(func):
    def retry_when_fail(*args, **kwargs):
        retry_times = 5
        for times in range(retry_times):
            if not func(*args, **kwargs):
                if times != (retry_times - 1):
                    print("run %s failed,retry times %s" % (func.__name__, times))
                    time.sleep(3)
                    continue
                else:
                    print("run %s failed.exit" % func.__name__)
                    raise
            print("run %s success." % func.__name__)
            break
        return func(*args, **kwargs)

    return retry_when_fail

@retry
def download_tool(ret):
    #download sth
    return ret

if __name__ == '__main__':
    #ret = True
    ret = False
    download_tool(ret)

ret = True和False的运行演示如下

带参数的装饰器

前面我们用到的@wraps(func)就是带参数的装饰器，允许传入func参数，下面我们定义一个类似的装饰器。在前面日志的基础上允许添加日志级别的参数。

import logging
logging.basicConfig(level=logging.INFO)

def logit(level="info"):
    def decorator(func):
        def with_called_logging(*args, **kwargs):
            if level == "info":
                logging.info(func.__name__ + " was called")
            elif level == "warn":
                logging.warn(func.__name__ + " was called")
            return func(*args, **kwargs)
        return with_called_logging

    return decorator

@logit()
def square_func(x):
    return x * x

@logit("warn")
def add_func(x, y):
    return x + y

s_result = square_func(4)
a_result = add_func(2, 3)
print(s_result, a_result)

这里相比之前多封装了一层函数，可以理解为含有参数的闭包。
由于logit()=decorator
所以@logit()、@logit(“warn”)都等于@decorator，日志级别相关参数通过闭包的形式传入。

类装饰器

装饰器不仅可以是函数，还可以是类，相比函数装饰器，类装饰器具有灵活度大、高内聚、封装性等优点。使用类装饰器主要依靠类的__call__方法，当使用 @ 形式将装饰器附加到函数上时，就会调用此方法。

import logging
logging.basicConfig(level=logging.INFO)

def logit(level="info"):
    class decorator(object):
        def __init__(self, func, level = "info"):
            self.level = level
            self._func = func
        def __call__(self,*args):
            if self.level == "info":
                logging.info(self._func.__name__ + " was called")
            elif self.level == "warn":
                logging.warn(self._func.__name__ + " was called")
            return self._func(*args)
    return decorator

@logit()
def square_func(x):
    return x * x

@logit(level = "warn")
def add_func(x, y):
    return x + y

s_result = square_func(4)
a_result = add_func(2, 3)
print(s_result, a_result)

__call__方法（补充知识点）

创建类型的时候定义了__call__()方法，这个类型就是可调用的。

import logging
logging.basicConfig(level=logging.INFO)

class decorator(object):
    def __init__(self, func, level = "info"):
        self.level = level
        self._func = func
    def __call__(self,*args):
        if self.level == "info":
            logging.info(self._func.__name__ + " was called")
        elif self.level == "warn":
            logging.warn(self._func.__name__ + " was called")
        return self._func(*args)

#@decorator
def square_func(x):
    return x * x

square_func = decorator(square_func)
print(type(square_func))
s_result = square_func(4)           #类的实例可以直接调用
print(s_result)

不含__call__方法的类型举例：

import logging
logging.basicConfig(level=logging.INFO)

class decorator(object):
    def __init__(self, func, level = "info"):
        self.level = level
        self._func = func
    def not_call(self,*args):
        if self.level == "info":
            logging.info(self._func.__name__ + " was called")
        elif self.level == "warn":
            logging.warn(self._func.__name__ + " was called")
        return self._func(*args)

#@decorator
def square_func(x):
    return x * x

square_func = decorator(square_func)
print(type(square_func))
try:
    s_result = square_func(4)
except TypeError as e:
    print(e)
s_result = square_func.not_call(4)
print(s_result)