带参数的全类型 Python 装饰器

这篇短文中显示的代码取自我的小型开源项目按合同设计，它提供了一个类型化的装饰器。装饰器是一个非常有用的概念，你肯定会在网上找到很多关于它们的介绍。简单说，它们允许在每次调用装饰函数时（之前和之后）执行代码。通过这种方式，你可以修改函数参数或返回值、测量执行时间、添加日志记录、执行执行时类型检查等等。请注意，装饰器也可以为类编写，提供另一种元编程方法（例如在 attrs 包中完成）

在最简单的形式中，装饰器的定义类似于以下代码：

def my_first_decorator(func):
   def wrapped(*args, **kwargs):
       # do something before
       result = func(*args, **kwargs)
       # do something after
       return result
   return wrapped
@my_first_decorator
def func(a):
   return a

如上代码，因为当定义了被包装的嵌套函数时，它的周围变量可以在函数内访问并保存在内存中，只要该函数在某处使用（这在函数式编程语言中称为闭包）。

很简单， 但是这有一些缺点。最大的问题是修饰函数会丢失它的之前的函数名字（你可以用inspect.signature看到这个），它的文档字符串，甚至它的名字， 这些是源代码文档工具（例如 sphinx）的问题，但可以使用标准库中的 functools.wraps 装饰器轻松解决：

from functools import wraps
from typing import Any, Callable, TypeVar, ParamSpec
P = ParamSpec("P") # 需要python >= 3.10
R = TypeVar("R")
def my_second_decorator(func: Callable[P, R]) -> Callable[P, R]:
   @wraps(func)
   def wrapped(*args: Any, **kwargs: Any) -> R:
       # do something before
       result = func(*args, **kwargs)
       # do something after
       return result
   return wrapped
@my_second_decorator
def func2(a: int) -> int:
   """Does nothing"""
   return a
print(func2.__name__)
# 'func2'
print(func2.__doc__)
# 'Does nothing'

在这个例子中，我已经添加了类型注释，注释和类型提示是对 Python 所做的最重要的补充。更好的可读性、IDE 中的代码完成以及更大代码库的可维护性只是其中的几个例子。上面的代码应该已经涵盖了大多数用例，但无法参数化装饰器。考虑编写一个装饰器来记录函数的执行时间，但前提是它超过了一定的秒数。这个数量应该可以为每个装饰函数单独配置。如果没有指定，则应使用默认值，并且应使用不带括号的装饰器，以便更易于使用：

@time(threshold=2)
def func1(a):
  ...
# No paranthesis when using default threshold
@time
def func2(b):
  ...

如果你可以在第二种情况下使用括号，或者根本不提供参数的默认值，那么这个秘诀就足够了：

from functools import wraps
from typing import Any, Callable, TypeVar, ParamSpec
P = ParamSpec("P") # 需要python >= 3.10
R = TypeVar("R")
def my_third_decorator(threshold: int = 1) -> Callable[[Callable[P, R]], Callable[P, R]]:
   def decorator(func: Callable[P, R]) -> Callable[P, R]:
       @wraps(func)
       def wrapper(*args: Any, **kwargs: Any) -> R:
           # do something before you can use `threshold`
           result = func(*args, **kwargs)
           # do something after
           return result
       return wrapper
   return decorator
@my_third_decorator(threshold=2)
def func3a(a: int) -> None:
  ...
# works
@my_third_decorator()
def func3b(a: int) -> None:
  ...
# Does not work!
@my_third_decorator
def func3c(a: int) -> None:
  ...

为了涵盖第三种情况，有一些包，即 wraps 和 decorator，它们实际上可以做的不仅仅是添加可选参数。虽然质量非常高，但它们引入了相当多的额外复杂性。使用 wrapt-decorated 函数，在远程集群上运行函数时，我进一步遇到了序列化问题。据我所知，两者都没有完全键入，因此静态类型检查器/ linter（例如 mypy）在严格模式下失败。

当我在自己的包上工作并决定编写自己的解决方案时，必须解决这些问题。它变成了一种可以轻松重用但很难转换为库的模式。

它使用标准库的重载装饰器。这样，可以指定相同的装饰器与我们的无参数一起使用。除此之外，它是上面两个片段的组合。这种方法的一个缺点是所有参数都需要作为关键字参数给出（这毕竟增加了可读性）

from typing import Callable, TypeVar, ParamSpec
from functools import partial, wraps
P = ParamSpec("P") # requires python >= 3.10
R = TypeVar("R
@overload
def typed_decorator(func: Callable[P, R]) -> Callable[P, R]:
  ...
@overload
def typed_decorator(*, first: str = "x", second: bool = True) -> Callable[[Callable[P, R]], Callable[P, R]]:
  ...
def typed_decorator(
   func: Optional[Callable[P, R]] = None, *, first: str = "x", second: bool = True
) -> Union[Callable[[Callable[P, R]], Callable[P, R]], Callable[P, R]]:
   """
  Describe what the decorator is supposed to do!
  Parameters
  ----------
  first : str, optional
      First argument, by default "x".
      This is a keyword-only argument!
  second : bool, optional
      Second argument, by default True.
      This is a keyword-only argument!
  """
   def wrapper(func: Callable[P, R], *args: Any, **kw: Any) -> R:
       """The actual logic"""
       # Do something with first and second and produce a `result` of type `R`
       return result
   # Without arguments `func` is passed directly to the decorator
   if func is not None:
       if not callable(func):
           raise TypeError("Not a callable. Did you use a non-keyword argument?")
       return wraps(func)(partial(wrapper, func))
   # With arguments, we need to return a function that accepts the function
   def decorator(func: Callable[P, R]) -> Callable[P, R]:
       return wraps(func)(partial(wrapper, func))
   return decorator

稍后，我们可以分别使用我们的不带参数的装饰器

@typed_decorator
def spam(a: int) -> int:
   return a
@typed_decorator(first = "y
def eggs(a: int) -> int:
   return a

这种模式肯定有一些开销，但收益大于成本。

原文：​ ​http://lemonfold.io/posts/2022/dbc/typed_decorator/​ ​