使用 arguments 提示 Python 上下文管理器的类型

Question

我有以下上下文管理器：

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

现在我有 function

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)的上下文管理器，或具有相同 arguments 的任何等效上下文管理器？

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

Answer 1

所以我在这里做一些假设来提供一个完整的答案。 如果其中任何一个不正确，您将不得不相应地调整代码。 我假设：

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

这里要理解的重要一件事是，您装饰的 function 本身不是上下文管理器。 它是一个工厂function，它返回一个上下文管理器，即一个 object，它的__enter__和__exit__方法定义为在with语句中使用。 （见文档）

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

在我看来，第一个在技术上是“最正确”的，但在您的情况以及许多其他情况下也可能是矫枉过正，首先定义您自己的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 pass（在--strict模式下）没有问题。

此处需要协议的原因是因为您的上下文管理器工厂中允许使用任意关键字 arguments **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工厂接受什么 arguments 的信息，因此您的 IDE 不会抱怨这样做：

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

您选择哪一个选项显然取决于您，并且取决于您想要的精确度。

希望这可以帮助。