Python学习：带参数的装饰器

问题

前面写了一个统计函数执行时间的装饰器，默认使用print函数打印执行日志。日志打印，不能假设一定会用print，能否使用自定义的日志打印，比如logger

方案

先看前面提供的exec_time装饰器的代码：

 

import time
from functools import wraps

def exec_time():
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            beg = time.time() * 1000
            ret = func(*args, **kwargs)
            end = time.time() * 1000
            print('func:{%s} exec time is:{%.5f} ms' % (func.__name__, end - beg))
            return ret

        return wrapper

    return decorator

补充说明下，这里使用了语言内建的装饰器@wraps，@wraps的作用是来保留函数元数据。不然，函数信息会丢失。

我们默认使用print打印，调用者也可以指定其他log打印的方式。实现代码如下：

 

def exec_time(log=None):
    """
    功能描述：定义嵌套函数，用来打印出装饰的函数的执行时间
    :param func: func参数（自动）
    :param log: 日志打印函数类型
    :return: 返回内部wrapper函数
    """
    logger = log

    def decorate(func):

        @wraps(func)
        def wrapper(*args, **kwargs):
            """
            功能描述：定义开始时间和结束时间，将func夹在中间执行，取得其返回值
            :param args:
            :param kwargs:
            :return: 被装饰函数的实际返回值
            """
            start = time.time() * 1000
            func_ret = func(*args, **kwargs)
            end = time.time() * 1000
            if logger is None:
                print('func:{%s} exec time is:{%.5f} ms' % (func.__name__, end - start))
            else:
                logger.debug('func:{%s} exec time is:{%.5f} ms' % (func.__name__, end - start))
            return func_ret

        return wrapper

    # 返回嵌套的函数
    return decorate

测试代码：

 

import time
from common.exec_time import exec_time

class MyLogger:
    def debug(self, str):
        print('mylogger:debug:' + str)

class TestExecTime:
    def test_exec_time_01(self):
        @exec_time()
        def sleep_1s():
            time.sleep(1)

        sleep_1s()

    def test_exec_time_02(self):
        @exec_time(MyLogger())
        def sleep_2s():
            time.sleep(2)

        sleep_2s()

输出：

 

platform darwin -- Python 3.8.2, pytest-6.1.2, py-1.9.0, pluggy-0.13.1 -- /usr/local/bin/python3.8
cachedir: .pytest_cache
rootdir: /Users/xxx/data/code/github/learn-python/test
collecting ... collected 2 items

test_ExecTime.py::TestExecTime::test_exec_time_01 
test_ExecTime.py::TestExecTime::test_exec_time_02 

============================== 2 passed in 3.02s ===============================

Process finished with exit code 0
PASSED                 [ 50%]func:{sleep_1s} exec time is:{1000.36304} ms
PASSED                 [100%]mylogger:debug:func:{sleep_2s} exec time is:{2002.80005} ms

讨论

开发库的时候，库的接口设计假设/依赖越少越好，一方面是考虑通用性，另外一方面也会使得库具备较好的扩展性。因为一旦接口固定，后期修改成本就高得多了。