Python 的上下文管理器协议只有两种方法，有直接的类型。 但是当涉及到为上下文管理器添加精确的类型提示时，我们仍然需要结合几种类型特征。 让我们看看我们如何为两种不同的方式制作的上下文管理器做到这一点。
@contextmanager 提示
@contextmanager 创建上下文管理器的最简单的方法是使用 contextlib 的装饰器。当添加类型提示时，这仍然是最简单的方法，因为我们只需要键入底层生成器。 例如。
from contextlib import contextmanager
from collections.abc import Iterator
@contextmanager
def my_context_manager() -> Iterator[None]:
yield
对于只产生值的简单生成器，我们可以用Iterator 来指定它们的类型。 注意：只有在 Python 3.9 上支持使用collections.abc.Iterator ；在旧版本上，我们需要导入typing.Iterator 来代替。
对于返回值的上下文管理器，我们可以将None 换成值的类型，比如说。
from contextlib import contextmanager
from collections.abc import Iterator
@contextmanager
def dice_roll() -> Iterator[int]:
yield 4
对于抑制异常的上下文管理器，我们将改为使用Generator ，这样我们就可以指定返回类型为bool ，例如。
from contextlib import contextmanager
from collections.abc import Generator
@contextmanager
def ignore_bad_math() -> Generator[None, None, bool]:
try:
yield
except ZeroDivisionError:
pass
return True
with ignore_bad_math():
1 / 0
同样，只有Python 3.9支持使用collections.abc.Generator ；在旧版本中，我们需要使用typing.Generator 。
基于类的上下文管理器
我们可以将更复杂的上下文管理器创建为类，这时类型提示需要做更多工作。 最简单的定义是这样的。
from __future__ import annotations
from types import TracebackType
class MyContextManager:
def __enter__(self) -> None:
pass
def __exit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> None:
pass
对于返回值的上下文管理器，我们会将__enter__'的返回类型从None 换成值的类型。
对于抑制异常的上下文管理器，我们将把__exit__'的返回类型改为bool 。
我们的__exit__ 方法的类型提示说。
exc_type是一个继承自 的类，或BaseExceptionNoneexc_val是 的一个实例（或一个子类），或BaseExceptionNoneexc_tb是一个回溯，或者None
这些都是真的，但是类型提示并不代表变量之间的关联性：它们要么都是设定的，要么都是None 。__exit__ ，永远不可能在只有某些值不是None 的情况下被调用。
如果我们只关心在with 语句中使用这个上下文管理器，我们可以根据需要在__exit__ 方法的主体中用类型缩小来处理这种相关性。但如果我们关心用户直接调用__exit__ ，我们可以达到@overload 。 让我们依次看看这两种技术。
窄化类型__exit__
为了处理里面的相关性，我们用if 和assert 添加一些类型缩小。如果我们想处理异常，我们可以使用这样的主体。
if exc_type is not None:
...
else:
...
Mypy可以使用if 语句来推断两个块中exc_type 的类型：在if 块中是type[BaseException] ，在else 块中是None 。但是由于Mypy不知道变量之间的相关性，它不能缩小exc_val 或exc_tb 的类型。我们可以告诉用assert 语句帮助Mypy缩小类型。
if exc_type is not None:
assert exc_val is not None
assert exc_tb is not None
...
else:
assert exc_val is None
assert exc_tb is None
...
Mypy可以读取asserts，并在下面几行中确定变量的类型。
上面的例子包含了完整的assert 语句，但我们并不总是需要如此详尽。如果一个块不使用某个变量，我们就不需要在那里缩小它的类型。 例如，许多上下文管理器只使用异常值，我们可以不提及其他变量。
if exc_val is not None:
... # do something with only exc_val
else:
...
如果对你需要的类型缩小有疑问，可以用reveal_type() 或reveal_locals() 进行调试。
使用@overload ，用于__exit__
为了使__exit__'的参数之间的相关性对调用者可见，我们需要使用@typing.overload来列出接受的形式。 这需要几个额外的存根函数。
from __future__ import annotations
from typing import overload
from types import TracebackType
class MyContextManager:
def __enter__(self) -> None:
pass
@overload
def __exit__(self, exc_type: None, exc_val: None, exc_tb: None) -> None:
...
@overload
def __exit__(
self,
exc_type: type[BaseException],
exc_val: BaseException,
exc_tb: TracebackType,
) -> None:
...
def __exit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> None:
pass
前两个@overload-decorated函数为调用者声明允许的类型。我们拼出两种情况：要么所有的参数都是None ，要么所有的参数都被设置。
最后一个__exit__ 函数是实现，在这里我们需要结合重载类型。注意，由于我们必须使用联合体，在主体内部我们仍然需要像上面那样使用类型缩小。（Mypy在当前不能将@overload 信息传播到主体中）。
现在，如果调用者试图传递一个不完整的参数集，他们将得到一个类型错误。 例如，如果我们这样写一个调用。
MyContextManager().__exit__(ValueError, None, None)
那么Mypy会这样抱怨。
$ mypy --strict example.py
example.py:40: error: No overload variant of "__exit__" of "MyContextManager" matches argument types "Type[ValueError]", "None", "None"
example.py:40: note: Possible overload variants:
example.py:40: note: def __exit__(self, None, None, None) -> None
example.py:40: note: def __exit__(self, Type[BaseException], BaseException, TracebackType) -> None
Found 1 error in 1 file (checked 1 source file)
Fin
我希望这能给你一些参考。
-阿丹
