Add support for narrowing booleans to literals

Author: ethan-leba

mypy/checker.py

@@ -4536,7 +4536,8 @@ def refine_identity_comparison_expression(self,
 
         enum_name = None
         target = get_proper_type(target)
-        if isinstance(target, LiteralType) and target.is_enum_literal():
+        if (isinstance(target, LiteralType) and
+                (target.is_enum_literal() or target.fallback.type.fullname == "builtins.bool")):
             enum_name = target.fallback.type.fullname
 
         target_type = [TypeRange(target, is_upper_bound=False)]

mypy/typeops.py

@@ -693,26 +693,32 @@ 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, 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():
-            if not isinstance(symbol.node, Var):
-                continue
-            # Skip "_order_" and "__order__", since Enum will remove it
-            if name in ("_order_", "__order__"):
-                continue
-            new_items.append(LiteralType(name, typ))
-        # SymbolTables are really just dicts, and dicts are guaranteed to preserve
-        # insertion order only starting with Python 3.7. So, we sort these for older
-        # versions of Python to help make tests deterministic.
-        #
-        # We could probably skip the sort for Python 3.6 since people probably run mypy
-        # 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, contract_literals=False)
-    else:
-        return typ
+    elif isinstance(typ, Instance) and typ.type.fullname == target_fullname:
+        if typ.type.is_enum:
+            new_items = []
+            for name, symbol in typ.type.names.items():
+                if not isinstance(symbol.node, Var):
+                    continue
+                # Skip "_order_" and "__order__", since Enum will remove it
+                if name in ("_order_", "__order__"):
+                    continue
+                new_items.append(LiteralType(name, typ))
+            # SymbolTables are really just dicts, and dicts are guaranteed to preserve
+            # insertion order only starting with Python 3.7. So, we sort these for older
+            # versions of Python to help make tests deterministic.
+            #
+            # We could probably skip the sort for Python 3.6 since people probably run mypy
+            # 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, contract_literals=False)
+        elif typ.type.fullname == "builtins.bool":
+            return make_simplified_union(
+                [LiteralType(True, typ), LiteralType(False, typ)],
+                contract_literals=False
+            )
+
+    return typ
 
 
 def try_contracting_literals_in_union(types: List[ProperType]) -> List[ProperType]:

test-data/unit/check-narrowing.test

@@ -1027,5 +1027,27 @@ if str_or_bool_literal is not True and str_or_bool_literal is not False:
     reveal_type(str_or_bool_literal)  # N: Revealed type is "builtins.str"
 else:
     reveal_type(str_or_bool_literal)  # N: Revealed type is "Union[Literal[False], Literal[True]]"
+[builtins fixtures/primitives.pyi]
+
+[case testNarrowingBoolean]
+# flags: --strict-optional
+from typing import Optional
+from typing_extensions import Literal
+
+bool_val: bool
 
+if bool_val is not False:
+    reveal_type(bool_val)   # N: Revealed type is "Literal[True]"
+else:
+    reveal_type(bool_val)   # N: Revealed type is "Literal[False]"
+
+opt_bool_val: Optional[bool]
+
+if opt_bool_val is not None:
+    reveal_type(opt_bool_val)   # N: Revealed type is "builtins.bool"
+
+if opt_bool_val is not False:
+    reveal_type(opt_bool_val)   # N: Revealed type is "Union[Literal[True], None]"
+else:
+    reveal_type(opt_bool_val)   # N: Revealed type is "Literal[False]"
 [builtins fixtures/primitives.pyi]