Python mypy 类型推断涉及联合和列表

Question

大家好，我对mypy在处理联合类型和列表时的行为有点困惑。这里有一个简化的例子来帮助解释：

from typing import Union


class A:
    pass

class B:
    pass


def f(items, a: Union[type[A], type[B]]) -> Union[list[A], list[B]]:
    return [a() for x in items] # Incompatible return value type (got "list[Union[A, B]]", expected "Union[list[A], list[B]]")

我原本以为这个列表推导式会创建一个只包含或的同类列表。但mypy却推断这个列表是混合的，可能同时包含和类型的项目。我不明白为什么会这样，因为参数要么是，要么是，在创建列表的过程中并没有变化。

从技术上讲，其实你甚至不需要使用列表推导式（虽然在我们的用例中，我们是用几个类来处理列表项）。下面的代码也会出错：

def f(a: Union[type[A], type[B]]) -> Union[list[A], list[B]]:
    return [a()]

我对这个情况的理解是否有误？有没有更好的方法来编写这个函数，以确保输出的类型是我想要的（我确实希望这个函数返回几个同类列表中的一个）。

Answer 1

这个类型检查器不知道传入的是A还是B，所以它不能更具体地确定a()的类型，只能说是A|B。因此，列表的类型是list[A|B]，而不是list[A] | list[B]。

你应该定义一个带有限制类型变量的通用函数。

from typing import TypeVar


T = TypeVar('T', A, B)


def f(items, t: type[T]) -> list[T]:
    return [t() for x in items]

从类型检查器的角度来看，你现在定义了两个函数，一个类型是Callable[[Any, Type[A]], list[A]]，另一个类型是Callable[[Any, Type[B]], list[B]]，而不是一个类型是Callable[[Any, Type[A]|Type[B]], list[A]|list[B]]的单一函数。参数t的类型实际上会决定调用哪一个函数，并且应该进行类型检查：

f(something, A)  # call the Callable[[Any, Type[A]], list[A]] version
f(something, B)  # call the Callable[[Any, Type[B]], list[B]] version