如何在Python泛型collect方法中保留集合及元素的类型提示？

原来的实现中，collect方法的返回类型被声明为Collection[T]，导致类型检查器无法推断出具体的集合子类类型（比如set[int]）。要保留具体集合类型和元素类型，可以通过引入额外的泛型变量来实现：

from typing import Generic, TypeVar, Collection
from dataclasses import dataclass

T = TypeVar("T")
# 定义绑定到Collection[T]的泛型变量C，代表具体的集合子类
C = TypeVar("C", bound=Collection[T])

@dataclass
class Example(Generic[T]):
    data: list[T]
    
    def collect(self, collector: type[C]) -> C:
        return collector(self.data)

这样修改后，当调用Example([1, 2]).collect(set)时，类型检查器会正确推断出result的类型为set[int]：

# 现在result的类型提示为'set[int]'，符合预期
result = Example([1, 2]).collect(set)

原理说明 #

• 新增的C是一个绑定到Collection[T]的泛型变量，它会匹配传入的具体集合类型（比如set、list、tuple等）。
• 方法参数collector: type[C]明确告诉类型检查器：传入的收集器是C类型的构造器。
• 返回类型C直接关联到传入的收集器类型，从而保留了具体集合的类型信息，同时元素类型T也会被正确传递。

如果需要支持更灵活的收集器（比如不是类型而是可接受可迭代对象的函数），可以将参数类型改为Callable[[Iterable[T]], C]：

from typing import Callable, Iterable

# 修改后的collect方法
def collect(self, collector: Callable[[Iterable[T]], C]) -> C:
    return collector(self.data)

这样既可以传入set这类类型构造器，也可以传入自定义的收集函数：

def custom_collector(iterable: Iterable[int]) -> frozenset[int]:
    return frozenset(iterable)

# result类型提示为'frozenset[int]'
result = Example([1, 2]).collect(custom_collector)

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