用 Python 写一个带参数的装饰器

如何解决用 Python 写一个带参数的装饰器

假设一个装饰器改变了一个函数 to_lower() 的行为：

def transition(fn):
    def to_upper(text: str):
        print('the original behavior:' + fn(text))
        return text.upper()
    
    def to_capital(text: str):
        print('the original behavior:' + fn(text))
        return text.capitalize()

    return to_upper

@transition
def to_lower(text: str):
    return text.lower()

print(to_lower('AaBb'))
>>>
the original behavior:aabb
AABB

这很好用，如果我们将 transition 的 return 语句从 return to_upper 更改为 return to_capital，它将行为从 to_lower 更改为 to_capital，这使 AaBb 变为 Aabb

不是手动修改装饰器的return， 如果我们调用，我们可以使用类似 mode 的参数修改装饰器吗？ @transition(mode='to_upper')，它的作用是 return to_upper，当我们调用 @transition(mode='to_capital')，它用作装饰器的 return to_capital？

解决方法

一个非常基本的实现：

def transition(mode):
    def deco(f):
        if mode == "to_upper":
            def wrapper(text):
                return text.upper()
        elif mode == "to_capital":
            def wrapper(text):
                return text.capitalize()
        else:
            # implement different behavior for other use cases  
            pass
        return wrapper
    return deco

    
    
        
    
@transition(mode="to_capital")
def to_lower(text):
    return text.lower()


@transition(mode="to_upper")
def to_lower(text):
    return text.lower()

,

你只需要为装饰器参数添加一个额外的层。示例：

import functools

def transition(mode):
    def wrapped(fn):
        @functools.wraps(fn)
        def inner(*args,**kwargs):
            res = fn(*args,**kwargs)
            if mode == 'to_capital':
                return res.capitalize()
            elif mode == 'to_upper':
                return res.upper()
            else:
                raise ValueError("invalid mode")

        return inner

    return wrapped

可以应用为

@transition(mode='to_capital')
def to_lower(text: str):
    return text.lower()

,

您可以将参数设为可选，代码也更简洁。

from functools import wraps

def process(fn,mode):
    @wraps(fn)
    def newfn(text):
        if mode == "to_upper":
            return text.upper()
        else:
            return text.capitalize()
    return newfn

def transition(fn=None,/,*,mode="to_upper"):
    def wrap(fn):
        return process(fn,mode)
    if fn is None:
        return wrap
    return wrap(fn)

@transition(mode="upper")
def to_lower(text: str):
    return text.lower()

@transition(mode="capitalize")
def to_lower(text: str):
    return text.lower()

@transition
def to_lower(text: str):
    return text.lower()