输入带有参数的Python contextmanager

Question

我有以下上下文管理器：

@contextmanager
def timed_task(task_name: str, **context_info)
     pass

现在我有了这样的功能

def my_func(my_timed_task: Any):
     with my_timed_task("my_func_task", foo="bar"):
          pass

我如何键入提示my_func，以使它知道my_timed_task是timed_task(task_name: str, **kwargs)类型的上下文管理器，或者任何具有相同参数的等效上下文管理器？

我知道contextlib.AbstractContextManager，但我可以找到说明如何将其与contextmanager参数结合使用的文档，而不仅仅是无争议的上下文管理器。

Answer 1

因此，我在这里作了几个假设，以提供一个完整的答案。如果其中任何一个是不正确的，您将不得不相应地调整代码。我想：

• 您的timed_task和my_func都可以返回任何类型。
• 你的**context_info卡可以是任何类型的。

这里需要理解的一件事是，您的修饰函数本身是而不是上下文管理器。它是一个工厂函数，返回上下文管理器，即定义了用于with-statement中的__enter__和__exit__方法的对象。(见文档)

在手头的事情上，你基本上有两个选择。

在我看来，first --在技术上是“最正确的”，但在您的情况和许多其他情况下也可能是过分的，它首先定义了您自己的Protocol：

from contextlib import contextmanager, AbstractContextManager
from typing import Any, Protocol


class MyContextManagerFactory(Protocol):
    def __call__(self, task_name: str, **context_info: Any) -> AbstractContextManager[Any]: ...


@contextmanager
def timed_task(task_name: str, **context_info: Any) -> Any:
    pass


def my_func(my_timed_task: MyContextManagerFactory) -> Any:
    with my_timed_task("my_func_task", foo="bar"):
        pass

这有mypy传递(在--strict模式下)，没有问题。

这里之所以需要协议，是因为您在上下文管理器工厂中允许任意关键字参数**context_info，而且据我所知，目前不可能相应地指定Callable类型。

对您的类型如此迂腐的好处是，像PyCharm这样的IDE不仅会给您提供精确的提示，说明在调用my_timed_task时允许使用什么类型的参数，而且还可以给出完整的签名(包括参数名)，这是很好的。

您的第二个选项要简单得多，但也不太精确。您可以简单地将my_timed_task类型定义为可调用的、接受任何内容并返回上下文管理器的类型：

from contextlib import contextmanager, AbstractContextManager
from typing import Any, Callable


@contextmanager
def timed_task(task_name: str, **context_info: Any) -> Any:
    pass


def my_func(my_timed_task: Callable[..., AbstractContextManager[Any]]) -> Any:
    with my_timed_task("my_func_task", foo="bar"):
        pass

同时也让mypy感到高兴。但是，没有关于您的my_timed_task工厂接受哪些参数的信息，所以您的IDE不会抱怨这样做：

def my_func(my_timed_task: Callable[..., AbstractContextManager[Any]]) -> Any:
    with my_timed_task(1, 3, "a", True):
        pass

那些选择中的哪一个显然取决于你自己，并取决于你想要的精确程度。

希望这能有所帮助。