Simplify unions when erasing last known values (#12064)

When we erase last known values in an union with multiple
Instance types, make sure that the resulting union doesn't have
duplicate erased types. The duplicate items weren't incorrect as such,
but they could cause overly complex error messages and potentially
slow type checking performance.

This is one of the fixes extracted from #12054. Since some of the
changes may cause regressions, it's better to split the PR.

Work on #12051.

Co-authored-by: Nikita Sobolev <[email protected]>

Author: JukkaL

mypy/erasetype.py

@@ -1,10 +1,10 @@
-from typing import Optional, Container, Callable
+from typing import Optional, Container, Callable, List, Dict, cast
 
 from mypy.types import (
     Type, TypeVisitor, UnboundType, AnyType, NoneType, TypeVarId, Instance, TypeVarType,
     CallableType, TupleType, TypedDictType, UnionType, Overloaded, ErasedType, PartialType,
     DeletedType, TypeTranslator, UninhabitedType, TypeType, TypeOfAny, LiteralType, ProperType,
-    get_proper_type, TypeAliasType, ParamSpecType
+    get_proper_type, get_proper_types, TypeAliasType, ParamSpecType
 )
 from mypy.nodes import ARG_STAR, ARG_STAR2
 
@@ -161,3 +161,34 @@ def visit_type_alias_type(self, t: TypeAliasType) -> Type:
         # Type aliases can't contain literal values, because they are
         # always constructed as explicit types.
         return t
+
+    def visit_union_type(self, t: UnionType) -> Type:
+        new = cast(UnionType, super().visit_union_type(t))
+        # Erasure can result in many duplicate items; merge them.
+        # Call make_simplified_union only on lists of instance types
+        # that all have the same fullname, to avoid simplifying too
+        # much.
+        instances = [item for item in new.items
+                     if isinstance(get_proper_type(item), Instance)]
+        # Avoid merge in simple cases such as optional types.
+        if len(instances) > 1:
+            instances_by_name: Dict[str, List[Instance]] = {}
+            new_items = get_proper_types(new.items)
+            for item in new_items:
+                if isinstance(item, Instance) and not item.args:
+                    instances_by_name.setdefault(item.type.fullname, []).append(item)
+            merged: List[Type] = []
+            for item in new_items:
+                if isinstance(item, Instance) and not item.args:
+                    types = instances_by_name.get(item.type.fullname)
+                    if types is not None:
+                        if len(types) == 1:
+                            merged.append(item)
+                        else:
+                            from mypy.typeops import make_simplified_union
+                            merged.append(make_simplified_union(types))
+                            del instances_by_name[item.type.fullname]
+                else:
+                    merged.append(item)
+            return UnionType.make_union(merged)
+        return new

mypy/test/testtypes.py

@@ -3,15 +3,16 @@
 from typing import List, Tuple
 
 from mypy.test.helpers import Suite, assert_equal, assert_type, skip
-from mypy.erasetype import erase_type
+from mypy.erasetype import erase_type, remove_instance_last_known_values
 from mypy.expandtype import expand_type
 from mypy.join import join_types, join_simple
 from mypy.meet import meet_types, narrow_declared_type
 from mypy.sametypes import is_same_type
 from mypy.indirection import TypeIndirectionVisitor
 from mypy.types import (
     UnboundType, AnyType, CallableType, TupleType, TypeVarType, Type, Instance, NoneType,
-    Overloaded, TypeType, UnionType, UninhabitedType, TypeVarId, TypeOfAny, get_proper_type
+    Overloaded, TypeType, UnionType, UninhabitedType, TypeVarId, TypeOfAny, ProperType,
+    get_proper_type
 )
 from mypy.nodes import ARG_POS, ARG_OPT, ARG_STAR, ARG_STAR2, CONTRAVARIANT, INVARIANT, COVARIANT
 from mypy.subtypes import is_subtype, is_more_precise, is_proper_subtype
@@ -1092,3 +1093,46 @@ def assert_simple_is_same(self, s: Type, t: Type, expected: bool, strict: bool)
                          '({} == {}) is {{}} ({{}} expected)'.format(s, t))
             assert_equal(hash(s) == hash(t), expected,
                          '(hash({}) == hash({}) is {{}} ({{}} expected)'.format(s, t))
+
+
+class RemoveLastKnownValueSuite(Suite):
+    def setUp(self) -> None:
+        self.fx = TypeFixture()
+
+    def test_optional(self) -> None:
+        t = UnionType.make_union([self.fx.a, self.fx.nonet])
+        self.assert_union_result(t, [self.fx.a, self.fx.nonet])
+
+    def test_two_instances(self) -> None:
+        t = UnionType.make_union([self.fx.a, self.fx.b])
+        self.assert_union_result(t, [self.fx.a, self.fx.b])
+
+    def test_multiple_same_instances(self) -> None:
+        t = UnionType.make_union([self.fx.a, self.fx.a])
+        assert remove_instance_last_known_values(t) == self.fx.a
+        t = UnionType.make_union([self.fx.a, self.fx.a, self.fx.b])
+        self.assert_union_result(t, [self.fx.a, self.fx.b])
+        t = UnionType.make_union([self.fx.a, self.fx.nonet, self.fx.a, self.fx.b])
+        self.assert_union_result(t, [self.fx.a, self.fx.nonet, self.fx.b])
+
+    def test_single_last_known_value(self) -> None:
+        t = UnionType.make_union([self.fx.lit1_inst, self.fx.nonet])
+        self.assert_union_result(t, [self.fx.a, self.fx.nonet])
+
+    def test_last_known_values_with_merge(self) -> None:
+        t = UnionType.make_union([self.fx.lit1_inst, self.fx.lit2_inst, self.fx.lit4_inst])
+        assert remove_instance_last_known_values(t) == self.fx.a
+        t = UnionType.make_union([self.fx.lit1_inst,
+                                  self.fx.b,
+                                  self.fx.lit2_inst,
+                                  self.fx.lit4_inst])
+        self.assert_union_result(t, [self.fx.a, self.fx.b])
+
+    def test_generics(self) -> None:
+        t = UnionType.make_union([self.fx.ga, self.fx.gb])
+        self.assert_union_result(t, [self.fx.ga, self.fx.gb])
+
+    def assert_union_result(self, t: ProperType, expected: List[Type]) -> None:
+        t2 = remove_instance_last_known_values(t)
+        assert type(t2) is UnionType
+        assert t2.items == expected

mypy/test/typefixture.py

@@ -157,9 +157,11 @@ def make_type_var(name: str, id: int, values: List[Type], upper_bound: Type,
         self.lit1 = LiteralType(1, self.a)
         self.lit2 = LiteralType(2, self.a)
         self.lit3 = LiteralType("foo", self.d)
+        self.lit4 = LiteralType(4, self.a)
         self.lit1_inst = Instance(self.ai, [], last_known_value=self.lit1)
         self.lit2_inst = Instance(self.ai, [], last_known_value=self.lit2)
         self.lit3_inst = Instance(self.di, [], last_known_value=self.lit3)
+        self.lit4_inst = Instance(self.ai, [], last_known_value=self.lit4)
 
         self.type_a = TypeType.make_normalized(self.a)
         self.type_b = TypeType.make_normalized(self.b)

test-data/unit/check-enum.test

@@ -1868,3 +1868,17 @@ class WithOverload(enum.IntEnum):
 class SubWithOverload(WithOverload):  # Should pass
     pass
 [builtins fixtures/tuple.pyi]
+
+[case testEnumtValueUnionSimplification]
+from enum import IntEnum
+from typing import Any
+
+class C(IntEnum):
+    X = 0
+    Y = 1
+    Z = 2
+
+def f1(c: C) -> None:
+    x = {'x': c.value}
+    reveal_type(x)  # N: Revealed type is "builtins.dict[builtins.str*, builtins.int]"
+[builtins fixtures/dict.pyi]