编写为包装函数添加参数的装饰器

编写一个装饰器,在包装函数的调用签名中添加一个额外的参数。 但是,添加的参数不会干扰该函数的现有调用约定。

额外的参数可以通过关键字参数注入到调用签名中,如下:

from functools import wraps


def optional_debug(func):
    if 'debug' in inspect.getfullargspec(func).args:
        raise TypeError('debug argument already defined')

    @wraps
    def wrapper(*args, debug=False, **kwargs):
        if debug:
            print('Calling', func.__name__)
        return func(*args, **kwargs)
    return wrapper


@optional_debug
def spam(a, b, c):
    print(a, b, c)

>>> spam(1, 2, 3)
1 2 3
>>> spam(1, 2, 3, debug=True)
Calling spam
1 2 3

若被包裹的函数已经定义了额外增加的参数,则程序会崩溃。

聪明的人会发现,使用装饰器后,新的函数的签名并没有包含debug参数,如下:

import inspect


@optional_debug
def add(x, y):
    return x + y

>>> print(inspect.signature(add))
(x, y)

这种bug可以通过以下方式修复:

from functools import wraps
import inspect


def optional_debug(func):
    if 'debug' in inspect.getfullargspec(func).args:
        raise TypeError('debug argument already defined')

    @wraps(func)
    def wrapper(*args, debug=False, **kwargs):
        if debug:
            print('Calling', func.__name__)
        return func(*args, **kwargs)

    sig = inspect.signature(func)
    parms = list(sig.parameters.values())
    parms.append(inspect.Parameter('debug', inspect.Parameter.KEYWORD_ONLY, default=False))
    wrapper.__signature__ = sig.replace(parameters=parms)
    return wrapper


@optional_debug
def add(x, y):
    return x + y

>>> print(inspect.signature(add))
(x, y, *, debug=False)
posted @ 2020-02-04 21:14  Jeffrey_Yang  阅读(194)  评论(0编辑  收藏  举报