我可以使用 python 装饰器根据输入类型预处理输入和后处理输出吗？答案

【问题标题】：Can I use a python decorator to preprocess input and postprocess output based on an input type?我可以使用 python 装饰器根据输入类型预处理输入和后处理输出吗？
【发布时间】：2021-05-04 03:50:31
【问题描述】：

我的函数可以采用List 或str。返回类型应与输入类型匹配。例如，如果给定 str，则函数应返回 str。

这是一个玩具例子来说明这一点：

def swap_first_and_last(a: Union[List, str]) -> Union[List, str]:
    # Not in-place.
    STR = isinstance(a, str)
    a = list(a)

    a[0], a[-1] = a[-1], a[0]

    return "".join(a) if STR else a

这是一个实际的例子：

def next_permutation(a: Union[List, str]) -> Union[List, str]:
    """
    Not in-place.
    Returns `None` if there is no next permutation
    (i.e. if `a` is the last permutation).
    The type of the output is the same as the type of the input
    (i.e. str input -> str output).
    """
    STR = isinstance(a, str)
    a = list(a)

    N = len(a)
    i = next((i for i in reversed(range(N-1)) if a[i] < a[i + 1]), None)
    if i is None:
        return None
    j = next(j for j in reversed(range(i+1, N)) if a[j] >= a[i])

    a[i], a[j] = a[j], a[i]
    a[i + 1:] = reversed(a[i + 1:])

    return "".join(a) if STR else a

如您所见，只有几行专门用于处理str 输入与List 输入，即：

    # preprocess
    STR = isinstance(a, str)
    a = list(a)
    
    # main logic
    ...
    
    # postprocess
    return "".join(a) if STR else a

我可以使用装饰器来做这种轻微的预处理和后处理吗？

【问题讨论】：

不完全是您所追求的，但您知道singledispatch 模块中的singledispatch 装饰器吗？

标签： python overloading decorator typing

【解决方案1】：

是的，你可以使用这样的装饰器：

from typing import Union, List

def pre_and_post_processing(func):
    def inner(a: Union[List, str]) -> Union[List, str]:
        STR = isinstance(a, str)
        a = list(a)
        b = func(a)
        return "".join(b) if STR else b
    return inner   


@pre_and_post_processing
def swap_first_and_last(a: List) -> List:
    a[0], a[-1] = a[-1], a[0]
    return a


print(swap_first_and_last("asd"))  # -> dsa
print(swap_first_and_last(["asd", "ds", "sds"]))  # -> ['sds', 'ds', 'asd']

请注意，swap_first_and_last 函数现在获取并返回 List。

【讨论】：

谢谢！我喜欢这两种解决方案，但是这个比singledispatch 的优点是可以定义一个装饰器（即@pre_and_post_processing）并在多个功能上使用它（即swap_fist_and_last 和next_permutation。相比之下，它似乎使用singledispatch，必须为每个功能设置singledispatch。

【解决方案2】：

跟进@larsks 的建议，这是在@joseville 玩具示例中使用singledispatch 的方法：

例子：

from functools import singledispatch

@singledispatch
def swap_first_and_last(a: str) -> str:

    a = list(a)
    swap_first_and_last(a)
    return "".join(a)

@swap_first_and_last.register
def _(a: list) -> list:
    
    a[0], a[-1] = a[-1], a[0]
    return a

print(swap_first_and_last([1,2,3]))
print(swap_first_and_last("123"))

输出：

[3, 2, 1]
321

【讨论】：