在 Python 中编写带参数的装饰器答案

【问题标题】：Write a decorator with parameter in Python在 Python 中编写带参数的装饰器
【发布时间】：2021-01-08 11:48:48
【问题描述】：

假设一个装饰器改变了一个函数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 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 工作？

【问题讨论】：

这个答案非常全面，可能会有所帮助：stackoverflow.com/a/5929165/7013263

标签： python decorator

【解决方案1】：

一个非常基本的实现：


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()

【讨论】：

【解决方案2】：

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

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()

【讨论】：

【解决方案3】：

您可以使参数成为可选参数，并且代码更简洁。

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()

【讨论】：