diff --git a/mypy/constraints.py b/mypy/constraints.py index 26986fa1975a..6b61d66fbf3b 100644 --- a/mypy/constraints.py +++ b/mypy/constraints.py @@ -386,6 +386,21 @@ def _infer_constraints( res.extend(infer_constraints(t_item, actual, direction)) return res if direction == SUPERTYPE_OF and isinstance(actual, UnionType): + # Case when *both* template and actual are unions needs special-casing: + # * First, remove identical items that appear in both unions. + # * Then, restore any unions that where split in the loop below + # (see comment above about plain type variables for why this is important). + # See e.g. testOptionalUnionInferencePrecise for situations where this helps. + nested_type_vars = None + if isinstance(template, UnionType): + nested_type_vars = [it for it in template.items if isinstance(it, TypeVarType)] + template, actual = remove_common_items(template, actual) + if not isinstance(actual, UnionType): + if isinstance(actual, UninhabitedType): + # Empty after simplification: do not infer spurious constraints. + return [] + # Not a union after simplification, restart from the top. + return infer_constraints(template, actual, direction) res = [] for a_item in actual.items: # `orig_template` has to be preserved intact in case it's recursive. @@ -394,6 +409,8 @@ def _infer_constraints( if type_type_unwrapped: a_item = TypeType.make_normalized(a_item) res.extend(infer_constraints(orig_template, a_item, direction)) + if nested_type_vars: + res = restore_union(res, actual, nested_type_vars) return res # Now the potential subtype is known not to be a Union or a type @@ -1498,6 +1515,37 @@ def find_matching_overload_items( return res +def remove_common_items(s: UnionType, t: UnionType) -> tuple[ProperType, ProperType]: + """Remove all items that appear in both unions.""" + common = set(s.items) & set(t.items) + new_s = UnionType.make_union([it for it in s.items if it not in common]) + new_t = UnionType.make_union([it for it in t.items if it not in common]) + return get_proper_type(new_s), get_proper_type(new_t) + + +def restore_union( + constraints: list[Constraint], union: UnionType, type_vars: list[TypeVarType] +) -> list[Constraint]: + """Merge constrains against each item of a union to a single constraint against the union. + + For each type variable we check whether all constraints for this type variable match + T :> Item for every item in the union, then replace such constraint group with a single + constraint. This will avoid accidentally inferring join from union. + """ + union_set = {get_proper_type(it) for it in union.items} + to_restore = set() + for tv in type_vars: + relevant_cs = [c for c in constraints if c.origin_type_var == tv] + if not all(c.op == SUPERTYPE_OF for c in relevant_cs): + continue + if union_set == {get_proper_type(c.target) for c in relevant_cs}: + to_restore.add(tv) + + original = [c for c in constraints if c.origin_type_var not in to_restore] + restored = [Constraint(tv, SUPERTYPE_OF, union) for tv in to_restore] + return original + restored + + def get_tuple_fallback_from_unpack(unpack: UnpackType) -> TypeInfo: """Get builtins.tuple type from available types to construct homogeneous tuples.""" tp = get_proper_type(unpack.type) diff --git a/test-data/unit/check-inference.test b/test-data/unit/check-inference.test index 88de271dfa12..05a4d5ece3d7 100644 --- a/test-data/unit/check-inference.test +++ b/test-data/unit/check-inference.test @@ -4347,3 +4347,49 @@ def func() -> None: [out] main:1: error: --local-partial-types must be enabled in parallel mode main:3: error: Need type annotation for "x" (hint: "x: list[] = ...") + +[case testOptionalUnionInferencePrecise] +from typing import Optional, TypeVar, Union + +T = TypeVar("T") + +def assert_not_none(x: Optional[T]) -> T: ... + +x: Union[int, str, None] +reveal_type(assert_not_none(x)) # N: Revealed type is "builtins.int | builtins.str" +y: Union[int, str] +reveal_type(assert_not_none(y)) # N: Revealed type is "builtins.int | builtins.str" +n: None +reveal_type(assert_not_none(n)) # N: Revealed type is "Never" + +def filter_not_none(x: list[Optional[T]]) -> list[T]: ... + +lx: list[Union[int, str, None]] +reveal_type(filter_not_none(lx)) # N: Revealed type is "builtins.list[builtins.int | builtins.str]" +reveal_type(filter_not_none([1, "yes", None])) # N: Revealed type is "builtins.list[builtins.str | builtins.int]" +[builtins fixtures/list.pyi] + +[case testStrictlyWiderInvariantUnionContextNoBottom] +from typing import TypeVar, Union + +T = TypeVar("T") +def foo(x: list[T]) -> list[Union[int, str, T]]: ... + +y: list[Union[int, str]] +y = foo(y) +[builtins fixtures/list.pyi] + +[case testOptionalUnionInferencePreciseSubclass] +from typing import Generic, Optional, TypeVar, Union + +T = TypeVar("T") +T_co = TypeVar("T_co", covariant=True) + +class A(Generic[T_co]): ... +class B(Generic[T_co]): ... +class C(A[Union[int, str]], B[Union[int, str, None]]): ... + +def f(x: Union[A[T], B[Optional[T]]]) -> T: ... + +reveal_type(f(C())) # N: Revealed type is "builtins.int | builtins.str" +[builtins fixtures/list.pyi] diff --git a/test-data/unit/check-python312.test b/test-data/unit/check-python312.test index 89f242a00333..9d32689a1e98 100644 --- a/test-data/unit/check-python312.test +++ b/test-data/unit/check-python312.test @@ -2373,3 +2373,31 @@ if TYPE_CHECKING: class Group(TypedDict): params: "Tuple[Alias, ...]" [builtins fixtures/tuple.pyi] + +[case testUnwrapUnionNewSyntaxPrecise] +from typing import Sequence + +class A: ... +class B: ... +class C: ... +type ThingOrA[T] = T | A + +def foo[T](x: ThingOrA[T]) -> T: ... +obj: B | C +reveal_type(foo(obj)) # N: Revealed type is "__main__.B | __main__.C" +obj_a: A | B | C +reveal_type(foo(obj_a)) # N: Revealed type is "__main__.B | __main__.C" + +def bar[T](xs: list[ThingOrA[T]]) -> list[T]: ... +objs: list[B | C] +objs_a: list[A | B | C] +reveal_type(bar(objs_a)) # N: Revealed type is "builtins.list[__main__.B | __main__.C]" +bar(objs) # E: Argument 1 to "bar" has incompatible type "list[B | C]"; expected "list[B | C | A]" \ + # N: "list" is invariant -- see https://mypy.readthedocs.io/en/stable/common_issues.html#variance \ + # N: Consider using "Sequence" instead, which is covariant + +def baz[T](xs: Sequence[ThingOrA[T]]) -> list[T]: ... +reveal_type(baz(objs)) # N: Revealed type is "builtins.list[__main__.B | __main__.C]" +reveal_type(baz(objs_a)) # N: Revealed type is "builtins.list[__main__.B | __main__.C]" +[builtins fixtures/tuple.pyi] +[typing fixtures/typing-full.pyi]