为类级contextmanager装饰器添加类型提示以支持自动补全

核心思路 #

要做一个和@contextlib.contextmanager等效的类装饰器，同时解决Pylance和mypy的类型推断问题，得结合Protocol定义上下文管理器接口和绑定型TypeVar，把装饰的类和返回的上下文实例类型绑定起来。

完整实现代码 #

from typing import TypeVar, Generic, Protocol, Iterator, Any, cast

# 绑定到类的TypeVar，关联装饰的类和上下文实例
C = TypeVar("C", bound=type)
# 协变TypeVar，对应__enter__的返回值类型
T = TypeVar("T", covariant=True)

class ClassBasedContextManager(Protocol[T]):
    def __enter__(self) -> T: ...
    def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> bool | None: ...

def class_contextmanager(cls: C) -> ClassBasedContextManager[T]:
    """类装饰器：把类转成上下文管理器，被装饰类必须实现__iter__返回单个值"""
    class _ContextManager(Generic[T]):
        def __init__(self, *args: Any, **kwargs: Any) -> None:
            self.instance = cls(*args, **kwargs)
            self.iterator: Iterator[T] | None = None

        def __enter__(self) -> T:
            self.iterator = iter(self.instance)
            value = next(self.iterator)
            return value

        def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> bool | None:
            if self.iterator is not None:
                try:
                    next(self.iterator)
                except StopIteration:
                    pass
            return None

    # 强制类型转换，让检查器识别内部类是合法的上下文管理器
    return cast(ClassBasedContextManager[T], _ContextManager)

# 使用示例
@class_contextmanager
class Example:
    class Internal:
        def do_something(self) -> None:
            print("Doing something")

    def __iter__(self) -> Iterator[Internal]:
        print("Entering context")
        yield self.Internal()
        print("Exiting context")

# 测试类型推断
with Example() as e:
    e.do_something()  # Pylance能识别e是Example.Internal，自动补全正常

关键细节说明 #

• Protocol接口：ClassBasedContextManager明确了上下文管理器的方法签名，让类型检查器能识别其结构，避免把内部类误判成函数。
• TypeVar绑定：C确保装饰器只接收类作为输入；T协变关联__enter__的返回值，解决类型不匹配问题。
• 类型转换：用cast让类型检查器认可内部类_ContextManager符合协议，解决Pylance的类型推断错误。
• 类的要求：被装饰的类必须实现__iter__方法，返回一个只产生单个值的迭代器，逻辑和@contextlib.contextmanager的生成器一致。

兼容性验证 #

• Pylance：能正确推断Example()是上下文管理器，with块里的e会被识别为Example.Internal，自动补全正常。
• mypy：开启--strict模式后无类型错误，能正确校验e.do_something()的调用合法性。

内容的提问来源于stack exchange，提问作者ted