Repository: quora/pyanalyze
Pull request #497 (2022-03-11 03:46:04)

Support type narrowing on `bool()` calls


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 68a9de4..1cafb33 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Support type narrowing on `bool()` calls (#497)
 - Support context managers that may suppress exceptions (#496)
 - Fix type inference for `with` assignment targets on
   Python 3.7 and higher (#495)
diff --git a/pyanalyze/implementation.py b/pyanalyze/implementation.py
index 5d84432..a30b7a1 100644
--- a/pyanalyze/implementation.py
+++ b/pyanalyze/implementation.py
@@ -1244,6 +1244,21 @@ def _len_impl(ctx: CallContext) -> ImplReturn:
     return ImplReturn(len_of_value(ctx.vars["obj"]), constraint)
 
 
+def _bool_impl(ctx: CallContext) -> Value:
+    if ctx.vars["o"] is _NO_ARG_SENTINEL:
+        return KnownValue(False)
+
+    # Maybe we should check boolability here too? But it seems fair to
+    # believe the author if they explicitly wrote bool().
+    varname = ctx.varname_for_arg("o")
+    if varname is None:
+        return TypedValue(bool)
+    constraint = Constraint(
+        varname, ConstraintType.is_truthy, positive=True, value=None
+    )
+    return annotate_with_constraint(TypedValue(bool), constraint)
+
+
 _POS_ONLY = ParameterKind.POSITIONAL_ONLY
 _ENCODING_PARAMETER = SigParameter(
     "encoding", annotation=TypedValue(str), default=KnownValue("")
@@ -1662,6 +1677,15 @@ def get_default_argspecs() -> Dict[object, Signature]:
             callable=len,
             impl=_len_impl,
         ),
+        Signature.make(
+            [
+                SigParameter(
+                    "o", ParameterKind.POSITIONAL_ONLY, default=_NO_ARG_SENTINEL
+                )
+            ],
+            callable=bool,
+            impl=_bool_impl,
+        ),
         # The overloaded annotation in typeshed causes a couple of problems:
         # - sorted(Sequence[A] | Sequence[B]) turns into List[A | B] instead of List[A] | List[B]
         #   pyright agrees with pyanalyze here but mypy somehow infers the second type.


Test patch:
diff --git a/pyanalyze/test_implementation.py b/pyanalyze/test_implementation.py
index 873bdc3..b5a0508 100644
--- a/pyanalyze/test_implementation.py
+++ b/pyanalyze/test_implementation.py
@@ -317,6 +317,24 @@ class TestLen(TestNameCheckVisitorBase):
             len(3)  # E: incompatible_argument
 
 
+class TestBool(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test_return_value(self):
+        def capybara(x):
+            assert_is_value(bool(), KnownValue(False))
+            assert_is_value(bool(x + 1), TypedValue(bool))
+
+    @assert_passes()
+    def test_constraint(self):
+        from typing import Optional
+
+        def capybara(x: Optional[int]):
+            if bool(x):
+                assert_is_value(x, TypedValue(int))
+            else:
+                assert_is_value(x, TypedValue(int) | KnownValue(None))
+
+
 class TestCast(TestNameCheckVisitorBase):
     @assert_passes()
     def test(self):
diff --git a/pyanalyze/test_stacked_scopes.py b/pyanalyze/test_stacked_scopes.py
index ecb337e..7a8ffee 100644
--- a/pyanalyze/test_stacked_scopes.py
+++ b/pyanalyze/test_stacked_scopes.py
@@ -1214,7 +1214,7 @@ class TestConstraints(TestNameCheckVisitorBase):
             if y is True:
                 assert_is_value(y, KnownValue(True))
                 y = bool(x)
-                assert_is_value(y, TypedValue(bool))
+                assert_is_value(y, TypedValue(bool), skip_annotated=True)
 
     @assert_passes()
     def test_constraint_on_arg_type(self):