Add some more inference (#35)

JelleZijlstra · web-flow · commit afb73ccbccb0 · 2022-08-21T15:27:23.000-07:00
Fixes #28
diff --git a/README.md b/README.md
@@ -82,6 +82,12 @@ Unreleased
 - Add `--pyanalyze-report`
 - Do not add `None` return types to methods marked with `@abstractmethod` and
   to methods in stub files
+- Improve type inference:
+  - `"string" % ...` is always `str`
+  - `b"bytes" % ...` is always `bytes`
+  - An `and` or `or` operator where left and right sides are of the same type
+    returns that type
+  - `is`, `is not`, `in`, and `not in` always return `bool`
 
 21.12.0 (December 21, 2021)
 
diff --git a/autotyping/autotyping.py b/autotyping/autotyping.py
@@ -560,6 +560,25 @@ def type_of_expression(expr: libcst.BaseExpression) -> Optional[Type[object]]:
         expr.operator, libcst.Not
     ):
         return bool
+    elif isinstance(expr, libcst.BinaryOperation):
+        left = type_of_expression(expr.left)
+        if left in (str, bytes) and isinstance(expr.operator, libcst.Modulo):
+            return left
+        return None
+    elif isinstance(expr, libcst.BooleanOperation):
+        left = type_of_expression(expr.left)
+        right = type_of_expression(expr.right)
+        # For AND and OR, if both types are the same, we can infer that type.
+        if left == right:
+            return left
+        else:
+            return None
+    elif isinstance(expr, libcst.Comparison):
+        types = {type(comp.operator) for comp in expr.comparisons}
+        # Only these are actually guaranteed to return bool
+        if types <= {libcst.In, libcst.Is, libcst.IsNot, libcst.NotIn}:
+            return bool
+        return None
     elif (
         isinstance(expr, libcst.Call)
         and isinstance(expr.func, libcst.Attribute)
diff --git a/tests/test_codemod.py b/tests/test_codemod.py
@@ -87,6 +87,15 @@ def return_not(x):
             def formatter(x):
                 return "{}".format(x)
 
+            def old_school_formatter(x):
+                return "%s" % x
+
+            def bytes_formatter(x):
+                return b"%s" % x
+
+            def boolean_return(x, y, z):
+                return x is y or x in z
+
             def baz() -> float:
                 return "not a float"
 
@@ -109,6 +118,15 @@ def return_not(x) -> bool:
             def formatter(x) -> str:
                 return "{}".format(x)
 
+            def old_school_formatter(x) -> str:
+                return "%s" % x
+
+            def bytes_formatter(x) -> bytes:
+                return b"%s" % x
+
+            def boolean_return(x, y, z) -> bool:
+                return x is y or x in z
+
             def baz() -> float:
                 return "not a float"
 
