Improve inference of union of generic types when one of the types is the generic type of the other

Author: roberfi

mypy/constraints.py

@@ -271,7 +271,11 @@ def infer_constraints_for_callable(
 
 
 def infer_constraints(
-    template: Type, actual: Type, direction: int, skip_neg_op: bool = False
+    template: Type,
+    actual: Type,
+    direction: int,
+    skip_neg_op: bool = False,
+    can_have_union_overlaping: bool = True,
 ) -> list[Constraint]:
     """Infer type constraints.
 
@@ -311,11 +315,15 @@ def infer_constraints(
         res = _infer_constraints(template, actual, direction, skip_neg_op)
         type_state.inferring.pop()
         return res
-    return _infer_constraints(template, actual, direction, skip_neg_op)
+    return _infer_constraints(template, actual, direction, skip_neg_op, can_have_union_overlaping)
 
 
 def _infer_constraints(
-    template: Type, actual: Type, direction: int, skip_neg_op: bool
+    template: Type,
+    actual: Type,
+    direction: int,
+    skip_neg_op: bool,
+    can_have_union_overlaping: bool = True,
 ) -> list[Constraint]:
     orig_template = template
     template = get_proper_type(template)
@@ -368,8 +376,41 @@ def _infer_constraints(
         return res
     if direction == SUPERTYPE_OF and isinstance(actual, UnionType):
         res = []
+
+        def _can_have_overlaping(_item: Type, _actual: UnionType) -> bool:
+            # There is a special overlaping case, where we have a Union of where two types
+            # are the same, but one of them contains the other.
+            # For example, we have Union[Sequence[T], Sequence[Sequence[T]]]
+            # In this case, only the second one can have overlaping because it contains the other.
+            # So, in case of list[list[int]], second one would be chosen.
+            if isinstance(p_item := get_proper_type(_item), Instance) and p_item.args:
+                other_items = [o_item for o_item in _actual.items if o_item is not a_item]
+
+                if len(other_items) == 1 and other_items[0] in p_item.args:
+                    return True
+
+                if len(other_items) > 1:
+                    union_args = [
+                        p_arg
+                        for arg in p_item.args
+                        if isinstance(p_arg := get_proper_type(arg), UnionType)
+                    ]
+
+                    for union_arg in union_args:
+                        if all(o_item in union_arg.items for o_item in other_items):
+                            return True
+
+            return False
+
         for a_item in actual.items:
-            res.extend(infer_constraints(orig_template, a_item, direction))
+            res.extend(
+                infer_constraints(
+                    orig_template,
+                    a_item,
+                    direction,
+                    can_have_union_overlaping=_can_have_overlaping(a_item, actual),
+                )
+            )
         return res
 
     # Now the potential subtype is known not to be a Union or a type
@@ -391,22 +432,28 @@ def _infer_constraints(
         # type variables indeterminate. This helps with some special
         # cases, though this isn't very principled.
 
-        def _is_item_being_overlaped_by_other(item: Type) -> bool:
-            # It returns true if the item is an argument of other item
+        def _is_item_overlaping_actual_type(_item: Type) -> bool:
+            # Overlaping occurs when we have a Union where two types are
+            # compatible and the more generic one is chosen.
+            # For example, in Union[T, Sequence[T]], we have to choose
+            # Sequence[T] if actual type is list[int].
+            # This returns true if the item is an argument of other item
             # that is subtype of the actual type
             return any(
-                isinstance(p_type := get_proper_type(item_to_compare), Instance)
-                and mypy.subtypes.is_subtype(actual, erase_typevars(p_type))
-                and item in p_type.args
+                isinstance(p_item_to_compare := get_proper_type(item_to_compare), Instance)
+                and mypy.subtypes.is_subtype(actual, erase_typevars(p_item_to_compare))
+                and _item in p_item_to_compare.args
                 for item_to_compare in template.items
-                if item is not item_to_compare
+                if _item is not item_to_compare
             )
 
         result = any_constraints(
             [
                 infer_constraints_if_possible(t_item, actual, direction)
                 for t_item in [
-                    item for item in template.items if not _is_item_being_overlaped_by_other(item)
+                    item
+                    for item in template.items
+                    if not (can_have_union_overlaping and _is_item_overlaping_actual_type(item))
                 ]
             ],
             eager=False,

test-data/unit/check-inference.test

@@ -885,7 +885,7 @@ i(b, a, b)
 i(a, b, b) # E: Argument 1 to "i" has incompatible type "List[int]"; expected "List[str]"
 [builtins fixtures/list.pyi]
 
-[case testUnionInferenceOfGenericClassAndItsGenericType]
+[case testInferenceOfUnionOfGenericClassAndItsGenericType]
 from typing import Generic, TypeVar, Union
 
 T = TypeVar('T')
@@ -907,6 +907,31 @@ reveal_type(result_2) # N: Revealed type is "__main__.GenericClass[builtins.str]
 
 [builtins fixtures/isinstance.pyi]
 
+[case testInferenceOfUnionOfSequenceOfAnyAndSequenceOfSequence]
+from typing import Sequence, Iterable, TypeVar, Union
+
+T = TypeVar("T")
+S = TypeVar("S")
+
+def sub_method(value: Union[S, Iterable[S]]) -> Iterable[S]:
+    pass
+
+def method(value: Union[Sequence[T], Sequence[Sequence[T]]]) -> None:
+    reveal_type(sub_method(value)) # N: Revealed type is "typing.Iterable[typing.Sequence[T`-1]]"
+
+[case testInferenceOfUnionOfUnionWithSequenceAndSequenceOfThatUnion]
+from typing import Sequence, Iterable, TypeVar, Union
+
+T = Union[str, Sequence[int]]
+S = TypeVar("S", bound=T)
+
+def sub_method(value: Union[S, Iterable[S]]) -> Iterable[S]:
+    pass
+
+def method(value: Union[T, Sequence[T]]) -> None:
+    reveal_type(sub_method(value)) # N: Revealed type is "typing.Iterable[Union[builtins.str, typing.Sequence[builtins.int]]]"
+
+
 [case testCallableListJoinInference]
 from typing import Any, Callable