Support for singleton types in unions

python · Michael0x2a · Oct 12, 2019 · Oct 11, 2019 · Oct 11, 2019 · Oct 11, 2019
commit 29b9b01b940aee531cc70218aad5f72e07a6bf9e
diff --git a/mypy/checker.py b/mypy/checker.py
@@ -4751,7 +4751,8 @@ def is_singleton_type(typ: Type) -> bool:
     'is_singleton_type(t)' returns True if and only if the expression 'a is b' is
     always true.
 
-    Currently, this returns True when given NoneTypes and enum LiteralTypes.
+    Currently, this returns True when given NoneTypes, enum LiteralTypes and
+    enum types with a single value.
 
     Note that other kinds of LiteralTypes cannot count as singleton types. For
     example, suppose we do 'a = 100000 + 1' and 'b = 100001'. It is not guaranteed
@@ -4761,7 +4762,8 @@ def is_singleton_type(typ: Type) -> bool:
     typ = get_proper_type(typ)
     # TODO: Also make this return True if the type is a bool LiteralType.
     # Also make this return True if the type corresponds to ... (ellipsis) or NotImplemented?
-    return isinstance(typ, NoneType) or (isinstance(typ, LiteralType) and typ.is_enum_literal())
+    return (isinstance(typ, NoneType) or (isinstance(typ, LiteralType) and typ.is_enum_literal())
+            or (isinstance(typ, Instance) and typ.type.is_enum and len(typ.type.names) == 1))
 
 
 def try_expanding_enum_to_union(typ: Type, target_fullname: str) -> ProperType:
@@ -4808,17 +4810,20 @@ class Status(Enum):
 
 
 def coerce_to_literal(typ: Type) -> ProperType:
-    """Recursively converts any Instances that have a last_known_value into the
-    corresponding LiteralType.
+    """Recursively converts any Instances that have a last_known_value or are
+    instances of enum types with a single value into the corresponding LiteralType.
     """
     typ = get_proper_type(typ)
     if isinstance(typ, UnionType):
         new_items = [coerce_to_literal(item) for item in typ.items]
         return make_simplified_union(new_items)
-    elif isinstance(typ, Instance) and typ.last_known_value:
-        return typ.last_known_value
-    else:
-        return typ
+    elif isinstance(typ, Instance):
+        if typ.last_known_value:
+            return typ.last_known_value
+        elif typ.type.is_enum and len(typ.type.names) == 1:
+            key, = typ.type.names
+            return LiteralType(value=key, fallback=typ)
+    return typ
 
 
 def has_bool_item(typ: ProperType) -> bool:

diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test
@@ -852,5 +852,29 @@ def process(response: Union[str, Reason] = '') -> str:
         # response can be only str, all other possible values exhausted
         reveal_type(response)  # N: Revealed type is 'builtins.str'
         return 'PROCESSED: ' + response
+[builtins fixtures/primitives.pyi]
+
+[case testEnumReachabilityPEP484Example3]
+# flags: --strict-optional
+from typing import Union
+from typing_extensions import Final
+from enum import Enum
 
+class Empty(Enum):
+    token = 0
+_empty = Empty.token
+
+def func(x: Union[int, None, Empty] = _empty) -> int:
+    boom = x + 42       # E: Unsupported left operand type for + ("None") \
+                        # E: Unsupported left operand type for + ("Empty") \
+                        # N: Left operand is of type "Union[int, None, Empty]"
+    if x is _empty:
+        reveal_type(x)  # N: Revealed type is 'Literal[__main__.Empty.token]'
+        return 0
+    elif x is None:
+        reveal_type(x)  # N: Revealed type is 'None'
+        return 1
+    else:  # At this point typechecker knows that x can only have type int
+        reveal_type(x)  # N: Revealed type is 'builtins.int'
+        return x + 2
 [builtins fixtures/primitives.pyi]