Recombine complete union of enum literals into original type (#9063) (#9097)

Closes #9063

Author: ethan-leba

mypy/join.py

@@ -310,8 +310,9 @@ def visit_literal_type(self, t: LiteralType) -> ProperType:
         if isinstance(self.s, LiteralType):
             if t == self.s:
                 return t
-            else:
-                return join_types(self.s.fallback, t.fallback)
+            if self.s.fallback.type.is_enum and t.fallback.type.is_enum:
+                return mypy.typeops.make_simplified_union([self.s, t])
+            return join_types(self.s.fallback, t.fallback)
         else:
             return join_types(self.s, t.fallback)
 

mypy/typeops.py

@@ -5,8 +5,9 @@
       since these may assume that MROs are ready.
 """
 
-from typing import cast, Optional, List, Sequence, Set, Iterable, TypeVar
+from typing import cast, Optional, List, Sequence, Set, Iterable, TypeVar, Dict, Tuple, Any
 from typing_extensions import Type as TypingType
+import itertools
 import sys
 
 from mypy.types import (
@@ -315,7 +316,8 @@ def callable_corresponding_argument(typ: CallableType,
 
 def make_simplified_union(items: Sequence[Type],
                           line: int = -1, column: int = -1,
-                          *, keep_erased: bool = False) -> ProperType:
+                          *, keep_erased: bool = False,
+                          contract_literals: bool = True) -> ProperType:
     """Build union type with redundant union items removed.
 
     If only a single item remains, this may return a non-union type.
@@ -377,6 +379,11 @@ def make_simplified_union(items: Sequence[Type],
                 items[i] = true_or_false(ti)
 
     simplified_set = [items[i] for i in range(len(items)) if i not in removed]
+
+    # If more than one literal exists in the union, try to simplify
+    if (contract_literals and sum(isinstance(item, LiteralType) for item in simplified_set) > 1):
+        simplified_set = try_contracting_literals_in_union(simplified_set)
+
     return UnionType.make_union(simplified_set, line, column)
 
 
@@ -684,7 +691,7 @@ class Status(Enum):
 
     if isinstance(typ, UnionType):
         items = [try_expanding_enum_to_union(item, target_fullname) for item in typ.items]
-        return make_simplified_union(items)
+        return make_simplified_union(items, contract_literals=False)
     elif isinstance(typ, Instance) and typ.type.is_enum and typ.type.fullname == target_fullname:
         new_items = []
         for name, symbol in typ.type.names.items():
@@ -702,11 +709,39 @@ class Status(Enum):
         # only using CPython, but we might as well for the sake of full correctness.
         if sys.version_info < (3, 7):
             new_items.sort(key=lambda lit: lit.value)
-        return make_simplified_union(new_items)
+        return make_simplified_union(new_items, contract_literals=False)
     else:
         return typ
 
 
+def try_contracting_literals_in_union(types: List[ProperType]) -> List[ProperType]:
+    """Contracts any literal types back into a sum type if possible.
+
+    Will replace the first instance of the literal with the sum type and
+    remove all others.
+
+    if we call `try_contracting_union(Literal[Color.RED, Color.BLUE, Color.YELLOW])`,
+    this function will return Color.
+    """
+    sum_types = {}  # type: Dict[str, Tuple[Set[Any], List[int]]]
+    marked_for_deletion = set()
+    for idx, typ in enumerate(types):
+        if isinstance(typ, LiteralType):
+            fullname = typ.fallback.type.fullname
+            if typ.fallback.type.is_enum:
+                if fullname not in sum_types:
+                    sum_types[fullname] = (set(get_enum_values(typ.fallback)), [])
+                literals, indexes = sum_types[fullname]
+                literals.discard(typ.value)
+                indexes.append(idx)
+                if not literals:
+                    first, *rest = indexes
+                    types[first] = typ.fallback
+                    marked_for_deletion |= set(rest)
+    return list(itertools.compress(types, [(i not in marked_for_deletion)
+                                           for i in range(len(types))]))
+
+
 def coerce_to_literal(typ: Type) -> Type:
     """Recursively converts any Instances that have a last_known_value or are
     instances of enum types with a single value into the corresponding LiteralType.

test-data/unit/check-enum.test

@@ -713,6 +713,7 @@ elif x is Foo.C:
     reveal_type(x)  # N: Revealed type is 'Literal[__main__.Foo.C]'
 else:
     reveal_type(x)  # No output here: this branch is unreachable
+reveal_type(x) # N: Revealed type is '__main__.Foo'
 
 if Foo.A is x:
     reveal_type(x)  # N: Revealed type is 'Literal[__main__.Foo.A]'
@@ -722,6 +723,7 @@ elif Foo.C is x:
     reveal_type(x)  # N: Revealed type is 'Literal[__main__.Foo.C]'
 else:
     reveal_type(x)  # No output here: this branch is unreachable
+reveal_type(x) # N: Revealed type is '__main__.Foo'
 
 y: Foo
 if y is Foo.A:
@@ -732,6 +734,7 @@ elif y is Foo.C:
     reveal_type(y)  # N: Revealed type is 'Literal[__main__.Foo.C]'
 else:
     reveal_type(y)  # No output here: this branch is unreachable
+reveal_type(y) # N: Revealed type is '__main__.Foo'
 
 if Foo.A is y:
     reveal_type(y)  # N: Revealed type is 'Literal[__main__.Foo.A]'
@@ -741,6 +744,7 @@ elif Foo.C is y:
     reveal_type(y)  # N: Revealed type is 'Literal[__main__.Foo.C]'
 else:
     reveal_type(y)  # No output here: this branch is unreachable
+reveal_type(y) # N: Revealed type is '__main__.Foo'
 [builtins fixtures/bool.pyi]
 
 [case testEnumReachabilityChecksWithOrdering]
@@ -815,12 +819,14 @@ if x is y:
 else:
     reveal_type(x)  # N: Revealed type is 'Union[Literal[__main__.Foo.B], Literal[__main__.Foo.C]]'
     reveal_type(y)  # N: Revealed type is 'Literal[__main__.Foo.A]'
+reveal_type(x) # N: Revealed type is '__main__.Foo'
 if y is x:
     reveal_type(x)  # N: Revealed type is 'Literal[__main__.Foo.A]'
     reveal_type(y)  # N: Revealed type is 'Literal[__main__.Foo.A]'
 else:
     reveal_type(x)  # N: Revealed type is 'Union[Literal[__main__.Foo.B], Literal[__main__.Foo.C]]'
     reveal_type(y)  # N: Revealed type is 'Literal[__main__.Foo.A]'
+reveal_type(x) # N: Revealed type is '__main__.Foo'
 
 if x is z:
     reveal_type(x)  # N: Revealed type is 'Literal[__main__.Foo.A]'
@@ -830,6 +836,7 @@ else:
     reveal_type(x)  # N: Revealed type is 'Union[Literal[__main__.Foo.B], Literal[__main__.Foo.C]]'
     reveal_type(z)  # N: Revealed type is 'Literal[__main__.Foo.A]?'
     accepts_foo_a(z)
+reveal_type(x) # N: Revealed type is '__main__.Foo'
 if z is x:
     reveal_type(x)  # N: Revealed type is 'Literal[__main__.Foo.A]'
     reveal_type(z)  # N: Revealed type is 'Literal[__main__.Foo.A]?'
@@ -838,6 +845,7 @@ else:
     reveal_type(x)  # N: Revealed type is 'Union[Literal[__main__.Foo.B], Literal[__main__.Foo.C]]'
     reveal_type(z)  # N: Revealed type is 'Literal[__main__.Foo.A]?'
     accepts_foo_a(z)
+reveal_type(x) # N: Revealed type is '__main__.Foo'
 
 if y is z:
     reveal_type(y)  # N: Revealed type is 'Literal[__main__.Foo.A]'
@@ -909,6 +917,7 @@ if x is Foo.A:
     reveal_type(x)  # N: Revealed type is 'Literal[__main__.Foo.A]'
 else:
     reveal_type(x)  # N: Revealed type is 'Union[Literal[__main__.Foo.B], Literal[__main__.Foo.C], None]'
+reveal_type(x) # N: Revealed type is 'Union[__main__.Foo, None]'
 [builtins fixtures/bool.pyi]
 
 [case testEnumReachabilityWithMultipleEnums]
@@ -928,18 +937,21 @@ if x1 is Foo.A:
     reveal_type(x1)  # N: Revealed type is 'Literal[__main__.Foo.A]'
 else:
     reveal_type(x1)  # N: Revealed type is 'Union[Literal[__main__.Foo.B], __main__.Bar]'
+reveal_type(x1) # N: Revealed type is 'Union[__main__.Foo, __main__.Bar]'
 
 x2: Union[Foo, Bar]
 if x2 is Bar.A:
     reveal_type(x2)  # N: Revealed type is 'Literal[__main__.Bar.A]'
 else:
     reveal_type(x2)  # N: Revealed type is 'Union[__main__.Foo, Literal[__main__.Bar.B]]'
+reveal_type(x2) # N: Revealed type is 'Union[__main__.Foo, __main__.Bar]'
 
 x3: Union[Foo, Bar]
 if x3 is Foo.A or x3 is Bar.A:
     reveal_type(x3)  # N: Revealed type is 'Union[Literal[__main__.Foo.A], Literal[__main__.Bar.A]]'
 else:
     reveal_type(x3)  # N: Revealed type is 'Union[Literal[__main__.Foo.B], Literal[__main__.Bar.B]]'
+reveal_type(x3) # N: Revealed type is 'Union[__main__.Foo, __main__.Bar]'
 
 [builtins fixtures/bool.pyi]
 
@@ -1299,3 +1311,25 @@ reveal_type(a._value_)  # N: Revealed type is 'Any'
 [builtins fixtures/__new__.pyi]
 [builtins fixtures/primitives.pyi]
 [typing fixtures/typing-medium.pyi]
+
+[case testEnumNarrowedToTwoLiterals]
+# Regression test: two literals of an enum would be joined
+# as the full type, regardless of the amount of elements
+# the enum contains.
+from enum import Enum
+from typing import Union
+from typing_extensions import Literal
+
+class Foo(Enum):
+    A = 1
+    B = 2
+    C = 3
+
+def f(x: Foo):
+    if x is Foo.A:
+        return x
+    if x is Foo.B:
+        pass
+    reveal_type(x) # N: Revealed type is 'Union[Literal[__main__.Foo.B], Literal[__main__.Foo.C]]'
+
+[builtins fixtures/bool.pyi]