如何根据参数值重载类实例成员的类型？

我最近了解到可以用Literal[True]和Literal[False]重载返回类型，现在要实现类似subprocess.Popen的接口，希望根据构造函数中text参数的值，将self.stdin的类型重载为IO[bytes]或IO[str]，使用pyright做静态类型检查。

尝试的代码及问题 #

我写了以下代码：

from typing import IO, Literal, Optional, overload


class MyPopen:
    @overload
    def __init__(self, text: Literal[False] = False):
        self.stdin: Optional[IO[bytes]]

    @overload
    def __init__(self, text: Literal[True] = True):
        self.stdin: Optional[IO[str]]

    def __init__(self, text: bool = False):
        self.stdin = None


pp = MyPopen(text=True)
assert pp.stdin
pp.stdin.write("text")    # 预期正常
pp.stdin.write(b"text")   # 预期报错
pp = MyPopen(text=False)
assert pp.stdin
pp.stdin.write("text")    # 预期报错
pp.stdin.write(b"text")   # 预期正常

但pyright推断两次pp.stdin都是IO[str]，还提示stdin声明被遮蔽，报错信息如下：

$ pyright /dev/stdin <1.py
/dev/stdin
  /dev/stdin:7:14 - error: Declaration "stdin" is obscured by a declaration of the same name (reportRedeclaration)
  /dev/stdin:20:1 - error: No overloads for "write" match the provided arguments (reportCallIssue)
  /dev/stdin:20:16 - error: Argument of type "Literal[b\"text\"]" cannot be assigned to parameter "__s" of type "str" in function "write"
    "Literal[b\"text\"]" is incompatible with "str" (reportArgumentType)
  /dev/stdin:24:1 - error: No overloads for "write" match the provided arguments (reportCallIssue)
  /dev/stdin:24:16 - error: Argument of type "Literal[b\"text\"]" cannot be assigned to parameter "__s" of type "str" in function "write"
    "Literal[b\"text\"]" is incompatible with "str" (reportArgumentType)
5 errors, 0 warnings, 0 informations 

参考案例：subprocess.Popen #

subprocess.Popen可以正常实现该功能：

import subprocess
pp = subprocess.Popen("", text=True)
assert pp.stdout
pp.stdout.write("text")   # 正常
pp.stdout.write(b"text")  # 报错
pp = subprocess.Popen("", text=False)
assert pp.stdout
pp.stdout.write("text")   # 报错
pp.stdout.write(b"text")  # 正常

我查看了CPython的subprocess源码，但没找到类型注解的位置。

解决方案 #

要实现根据构造参数重载实例成员类型，需要结合泛型（Generic）和类型变量（TypeVar），让类的类型与text参数的字面量类型绑定，具体实现如下：

正确代码示例 #

from typing import IO, Literal, Optional, overload, Generic, TypeVar

# 定义绑定到Literal[True]/Literal[False]的类型变量
TextType = TypeVar('TextType', Literal[True], Literal[False])

class MyPopen(Generic[TextType]):
    # 根据TextType的类型动态确定stdin的类型
    stdin: Optional[IO[str]] if TextType is Literal[True] else Optional[IO[bytes]]

    @overload
    def __init__(self, text: Literal[True] = ...) -> None:
        ...

    @overload
    def __init__(self, text: Literal[False] = ...) -> None:
        ...

    def __init__(self, text: bool = False) -> None:
        self.stdin = None

# 测试验证
pp = MyPopen(text=True)
assert pp.stdin
pp.stdin.write("text")    # pyright无报错
pp.stdin.write(b"text")   # pyright提示类型不兼容（符合预期）

pp = MyPopen(text=False)
assert pp.stdin
pp.stdin.write("text")    # pyright提示类型不兼容（符合预期）
pp.stdin.write(b"text")   # pyright无报错

原理说明 #

1. 类型变量绑定：通过TypeVar定义只能取Literal[True]或Literal[False]的类型变量TextType，将类的泛型类型与text参数的字面量类型关联。
2. 条件类型注解：利用Python的条件表达式，根据TextType的具体值，动态指定stdin的类型为IO[str]或IO[bytes]。
3. 重载构造函数：用@overload标注__init__的不同参数情况，让类型检查器能根据传入的text值推断出实例的具体类型。

关于subprocess的类型注解 #

subprocess的类型注解并不在CPython的源码中，而是存放在标准库的类型注解仓库里。它的实现逻辑和上面的示例一致：通过泛型绑定参数的字面量类型，动态推导流对象的类型。

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