Expand booleans during truthy checks and bool ops

Author: ethan-leba

mypy/checker.py

@@ -4318,7 +4318,7 @@ def has_no_custom_eq_checks(t: Type) -> bool:
         elif isinstance(node, RefExpr):
             # Restrict the type of the variable to True-ish/False-ish in the if and else branches
             # respectively
-            vartype = type_map[node]
+            vartype = try_expanding_sum_type_to_union(type_map[node], "builtins.bool")
             if_type: Type = true_only(vartype)
             else_type: Type = false_only(vartype)
             ref: Expression = node

mypy/checkexpr.py

@@ -64,8 +64,9 @@
     FunctionContext, FunctionSigContext,
 )
 from mypy.typeops import (
-    tuple_fallback, make_simplified_union, true_only, false_only, erase_to_union_or_bound,
-    function_type, callable_type, try_getting_str_literals, custom_special_method,
+    try_expanding_sum_type_to_union, tuple_fallback, make_simplified_union,
+    true_only, false_only, erase_to_union_or_bound, function_type,
+    callable_type, try_getting_str_literals, custom_special_method,
     is_literal_type_like,
 )
 import mypy.errorcodes as codes
@@ -2783,6 +2784,9 @@ def check_boolean_op(self, e: OpExpr, context: Context) -> Type:
         # '[1] or []' are inferred correctly.
         ctx = self.type_context[-1]
         left_type = self.accept(e.left, ctx)
+        expanded_left_type = try_expanding_sum_type_to_union(
+            self.accept(e.left, ctx), "builtins.bool"
+        )
 
         assert e.op in ('and', 'or')  # Checked by visit_op_expr
 
@@ -2817,7 +2821,7 @@ def check_boolean_op(self, e: OpExpr, context: Context) -> Type:
         # to be unreachable and therefore any errors found in the right branch
         # should be suppressed.
         with (self.msg.disable_errors() if right_map is None else nullcontext()):
-            right_type = self.analyze_cond_branch(right_map, e.right, left_type)
+            right_type = self.analyze_cond_branch(right_map, e.right, expanded_left_type)
 
         if right_map is None:
             # The boolean expression is statically known to be the left value

test-data/unit/check-dynamic-typing.test

@@ -159,9 +159,9 @@ a or d
 if int():
     c = a in d  # E: Incompatible types in assignment (expression has type "bool", variable has type "C")
 if int():
-    c = b and d # E: Incompatible types in assignment (expression has type "Union[bool, Any]", variable has type "C")
+    c = b and d # E: Incompatible types in assignment (expression has type "Union[Literal[False], Any]", variable has type "C")
 if int():
-    c = b or d  # E: Incompatible types in assignment (expression has type "Union[bool, Any]", variable has type "C")
+    c = b or d  # E: Incompatible types in assignment (expression has type "Union[Literal[True], Any]", variable has type "C")
 if int():
     b = a + d
 if int():

test-data/unit/check-expressions.test

@@ -316,11 +316,11 @@ if int():
 if int():
     b = b or b
 if int():
-    b = b and a # E: Incompatible types in assignment (expression has type "Union[bool, A]", variable has type "bool")
+    b = b and a # E: Incompatible types in assignment (expression has type "Union[Literal[False], A]", variable has type "bool")
 if int():
     b = a and b # E: Incompatible types in assignment (expression has type "Union[A, bool]", variable has type "bool")
 if int():
-    b = b or a  # E: Incompatible types in assignment (expression has type "Union[bool, A]", variable has type "bool")
+    b = b or a  # E: Incompatible types in assignment (expression has type "Union[Literal[True], A]", variable has type "bool")
 if int():
     b = a or b  # E: Incompatible types in assignment (expression has type "Union[A, bool]", variable has type "bool")
 class A: pass

test-data/unit/check-narrowing.test

@@ -1052,6 +1052,57 @@ else:
     reveal_type(opt_bool_val)   # N: Revealed type is "Literal[False]"
 [builtins fixtures/primitives.pyi]
 
+[case testNarrowingBooleanTruthiness]
+# flags: --strict-optional
+from typing import Optional
+from typing_extensions import Literal
+
+bool_val: bool
+
+if bool_val:
+    reveal_type(bool_val)   # N: Revealed type is "Literal[True]"
+else:
+    reveal_type(bool_val)   # N: Revealed type is "Literal[False]"
+reveal_type(bool_val)  # N: Revealed type is "builtins.bool"
+
+opt_bool_val: Optional[bool]
+
+if opt_bool_val:
+    reveal_type(opt_bool_val)   # N: Revealed type is "Literal[True]"
+else:
+    reveal_type(opt_bool_val)   # N: Revealed type is "Union[Literal[False], None]"
+reveal_type(opt_bool_val)   # N: Revealed type is "Union[builtins.bool, None]"
+[builtins fixtures/primitives.pyi]
+
+[case testNarrowingBooleanBoolOp]
+# flags: --strict-optional
+from typing import Optional
+from typing_extensions import Literal
+
+bool_a: bool
+bool_b: bool
+
+if bool_a and bool_b:
+   reveal_type(bool_a) # N: Revealed type is "Literal[True]"
+   reveal_type(bool_b) # N: Revealed type is "Literal[True]"
+else:
+   reveal_type(bool_a) # N: Revealed type is "builtins.bool"
+   reveal_type(bool_b) # N: Revealed type is "builtins.bool"
+
+if not bool_a or bool_b:
+   reveal_type(bool_a) # N: Revealed type is "builtins.bool"
+   reveal_type(bool_b) # N: Revealed type is "builtins.bool"
+else:
+   reveal_type(bool_a) # N: Revealed type is "Literal[True]"
+   reveal_type(bool_b) # N: Revealed type is "Literal[False]"
+
+if True and bool_b:
+   reveal_type(bool_b) # N: Revealed type is "Literal[True]"
+
+x = True and bool_b
+reveal_type(x) # N: Revealed type is "builtins.bool"
+[builtins fixtures/primitives.pyi]
+
 [case testNarrowingTypedDictUsingEnumLiteral]
 # flags: --python-version 3.6
 from typing import Union

test-data/unit/check-newsemanal.test

@@ -304,6 +304,7 @@ from a import x
 def f(): pass
 
 [targets a, b, a, a.y, b.f, __main__]
+[builtins fixtures/tuple.pyi]
 
 [case testNewAnalyzerRedefinitionAndDeferral1b]
 import a

test-data/unit/check-unreachable-code.test

@@ -533,11 +533,11 @@ f = (PY3 or PY2) and 's'
 g = (PY2 or PY3) or 's'
 h = (PY3 or PY2) or 's'
 reveal_type(a)  # N: Revealed type is "builtins.bool"
-reveal_type(b)  # N: Revealed type is "builtins.str"
-reveal_type(c)  # N: Revealed type is "builtins.str"
+reveal_type(b)  # N: Revealed type is "Literal['s']"
+reveal_type(c)  # N: Revealed type is "Literal['s']"
 reveal_type(d)  # N: Revealed type is "builtins.bool"
-reveal_type(e)  # N: Revealed type is "builtins.str"
-reveal_type(f)  # N: Revealed type is "builtins.str"
+reveal_type(e)  # N: Revealed type is "Literal['s']"
+reveal_type(f)  # N: Revealed type is "Literal['s']"
 reveal_type(g)  # N: Revealed type is "builtins.bool"
 reveal_type(h)  # N: Revealed type is "builtins.bool"
 [builtins fixtures/ops.pyi]

test-data/unit/typexport-basic.test

@@ -247,7 +247,7 @@ elif not a:
 [out]
 NameExpr(3) : builtins.bool
 IntExpr(4) : Literal[1]?
-NameExpr(5) : builtins.bool
+NameExpr(5) : Literal[False]
 UnaryExpr(5) : builtins.bool
 IntExpr(6) : Literal[1]?
 
@@ -259,7 +259,7 @@ while a:
 [builtins fixtures/bool.pyi]
 [out]
 NameExpr(3) : builtins.bool
-NameExpr(4) : builtins.bool
+NameExpr(4) : Literal[True]
 
 
 -- Simple type inference