如何关联Generic类型参数的可选性，确保两者匹配？

问题背景 #

你需要定义泛型类Worker，要求类型参数I和O的可选性必须一致：要么同时为Optional[...]类型，要么同时为非可选类型。但当前的实现无法限制这种关联性，允许创建一个为可选、另一个为非可选的错误实例（比如Worker[Optional[list], dict]()）。你目前通过创建两个独立类实现了类型限制，但不想重复代码，希望找到更简洁的方案。

原问题代码 #

from typing import Generic, TypeVar, Optional, Iterable, Mappable, abstractmethod

I = TypeVar("I", bound=Optional[Iterable])
O = TypeVar("O", bound=Optional[Mappable])

class Worker(Generic[I, O]):
    @abstractmethod
    def do_work(self, input: I) -> O:
        pass

# 合法实例
worker = Worker[list, dict]()
worker_with_optional = Worker[Optional[list], Optional[dict]]()
# 不希望允许的错误实例，但当前类型检查不报错
worker_bad_types = Worker[Optional[list], dict]()

当前替代方案（存在代码重复） #

from typing import Generic, TypeVar, Optional, Iterable, Mappable, abstractmethod

I = TypeVar("I", bound=Iterable)
O = TypeVar("O", bound=Mappable)

class Worker(Generic[I, O]):
    @abstractmethod
    def do_work(self, input: I) -> O:
        pass

class WorkerWithOptional(Generic[I, O]):
    @abstractmethod
    def do_work(self, input: Optional[I]) -> Optional[O]:
        pass

# 合法实例
worker = Worker[list, dict]()
worker_with_optional = WorkerWithOptional[list, dict]()
# 现在类型检查会报错，符合预期
worker_bad_types = Worker[Optional[int], str]()

解决方案 #

可以通过引入一个可选性标记类型参数，将I和O的可选性绑定在一起，从而避免重复类定义。具体实现如下：

实现代码 #

from typing import Generic, TypeVar, Optional, Iterable, Mappable, abstractmethod, Never

# 定义基础的非可选类型变量，约束为Iterable/Mappable
I_Base = TypeVar("I_Base", bound=Iterable)
O_Base = TypeVar("O_Base", bound=Mappable)
# 定义可选性标记，只能取True或False
IsOptional = TypeVar("IsOptional", bool)

class Worker(Generic[I_Base, O_Base, IsOptional]):
    @abstractmethod
    def do_work(
        self,
        input: I_Base | (Optional[I_Base] if IsOptional else Never),
    ) -> O_Base | (Optional[O_Base] if IsOptional else Never):
        pass

# 非可选版本：标记为False，input和返回值都是基础类型
worker = Worker[list, dict, False]()
# 可选版本：标记为True，input和返回值都是Optional包裹的基础类型
worker_with_optional = Worker[list, dict, True]()

# 以下两种错误场景都会触发类型检查报错：
# 1. 尝试传入已包裹Optional的基础类型（I_Base必须是未包裹的Iterable）
worker_bad_1 = Worker[Optional[list], dict, False]()
# 2. 标记为True但实现中返回非Optional类型（类型不匹配）
worker_bad_2 = Worker[list, dict, True]()

简化版本（用类型别名） #

from typing import Generic, TypeVar, Optional, Iterable, Mappable, abstractmethod, Never, TypeAlias

I_Base = TypeVar("I_Base", bound=Iterable)
O_Base = TypeVar("O_Base", bound=Mappable)
OptFlag = TypeVar("OptFlag", bool)

# 根据标记定义统一的输入输出类型别名
InputType: TypeAlias = I_Base | (Optional[I_Base] if OptFlag else Never)
OutputType: TypeAlias = O_Base | (Optional[O_Base] if OptFlag else Never)

class Worker(Generic[I_Base, O_Base, OptFlag]):
    @abstractmethod
    def do_work(self, input: InputType) -> OutputType:
        pass

这种方案的核心是通过一个统一的标记参数OptFlag控制I和O的可选性，确保两者始终保持一致，同时只需要维护一个类，彻底避免代码重复。

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