Simplify unions when erasing last known values

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]