Repository: quora/pyanalyze
Pull request #375 (2021-12-28 13:49:30)

Support the walrus operator


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 7a3a321..5922dfa 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Support the walrus (`:=`) operator (#375)
 - Initial support for proposed new "type evaluation"
   mechanism (#374)
 - Create command-line options for each config option (#373)
diff --git a/pyanalyze/ast_annotator.py b/pyanalyze/ast_annotator.py
index 00e5c4d..89d78da 100644
--- a/pyanalyze/ast_annotator.py
+++ b/pyanalyze/ast_annotator.py
@@ -9,6 +9,7 @@ attributes.
 """
 import ast
 import os
+import textwrap
 import traceback
 import types
 from typing import Optional, Type, Union
@@ -53,6 +54,7 @@ def annotate_code(
     :type verbose: bool
 
     """
+    code = textwrap.dedent(code)
     tree = ast.parse(code)
     try:
         mod = make_module(code)
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index 070dbdf..8e691a7 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -82,6 +82,7 @@ from .shared_options import Paths, ImportPaths, EnforceNoUnused
 from .reexport import ImplicitReexportTracker
 from .safe import safe_getattr, is_hashable, all_of_type
 from .stacked_scopes import (
+    EMPTY_ORIGIN,
     AbstractConstraint,
     Composite,
     CompositeIndex,
@@ -174,6 +175,11 @@ from .value import (
     unpack_values,
 )
 
+try:
+    from ast import NamedExpr
+except ImportError:
+    NamedExpr = Any  # 3.7 and lower
+
 T = TypeVar("T")
 AwaitableValue = GenericValue(collections.abc.Awaitable, [TypeVarValue(T)])
 KnownNone = KnownValue(None)
@@ -1198,7 +1204,7 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor, CanAssignContext):
         value: Value = AnyValue(AnySource.inference),
         *,
         private: bool = False,
-    ) -> Value:
+    ) -> Tuple[Value, VarnameOrigin]:
         current_scope = self.scopes.current_scope()
         scope_type = current_scope.scope_type
         if self.module is not None and scope_type == ScopeType.module_scope:
@@ -1208,12 +1214,12 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor, CanAssignContext):
                 )
             if varname in current_scope:
                 value, _ = current_scope.get_local(varname, node, self.state)
-                return value
+                return value, EMPTY_ORIGIN
         if scope_type == ScopeType.class_scope and isinstance(node, ast.AST):
             self._check_for_incompatible_overrides(varname, node, value)
             self._check_for_class_variable_redefinition(varname, node)
-        current_scope.set(varname, value, node, self.state)
-        return value
+        origin = current_scope.set(varname, value, node, self.state)
+        return value, origin
 
     def _check_for_incompatible_overrides(
         self, varname: str, node: ast.AST, value: Value
@@ -1465,7 +1471,8 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor, CanAssignContext):
         self._generic_visit_list(node.bases)
         self._generic_visit_list(node.keywords)
         value = self._visit_class_and_get_value(node)
-        return self._set_name_in_scope(node.name, node, value)
+        value, _ = self._set_name_in_scope(node.name, node, value)
+        return value
 
     def _get_class_object(self, node: ast.ClassDef) -> Value:
         if self.scopes.scope_type() == ScopeType.module_scope:
@@ -2230,9 +2237,10 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor, CanAssignContext):
                         for val in iterable_type:
                             self.visit_comprehension(generator, iterable_type=val)
                             with qcore.override(self, "in_comprehension_body", True):
-                                items.append(
-                                    KVPair(self.visit(node.key), self.visit(node.value))
-                                )
+                                # PEP 572 mandates that the key be evaluated first.
+                                key = self.visit(node.key)
+                                value = self.visit(node.value)
+                                items.append(KVPair(key, value))
                     finally:
                         self.node_context.contexts.pop()
                     return DictIncompleteValue(typ, items)
@@ -3448,6 +3456,20 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor, CanAssignContext):
 
     # Assignments
 
+    def visit_NamedExpr(self, node: NamedExpr) -> Value:
+        composite = self.composite_from_walrus(node)
+        return composite.value
+
+    def composite_from_walrus(self, node: NamedExpr) -> Composite:
+        rhs = self.visit(node.value)
+        with qcore.override(self, "being_assigned", rhs):
+            if self.in_comprehension_body:
+                ctx = self.scopes.ignore_topmost_scope()
+            else:
+                ctx = qcore.empty_context
+            with ctx:
+                return self.composite_from_node(node.target)
+
     def visit_Assign(self, node: ast.Assign) -> None:
         is_yield = isinstance(node.value, ast.Yield)
         value = self.visit(node.value)
@@ -3572,10 +3594,14 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor, CanAssignContext):
             else:
                 is_ann_assign = False
             value = self.being_assigned
+            origin = EMPTY_ORIGIN
             if not is_ann_assign:
                 self.yield_checker.record_assignment(node.id)
-                value = self._set_name_in_scope(node.id, node, value=value)
-            return Composite(value, VarnameWithOrigin(node.id), node)
+                value, origin = self._set_name_in_scope(node.id, node, value=value)
+            varname = VarnameWithOrigin(node.id, origin)
+            constraint = Constraint(varname, ConstraintType.is_truthy, True, None)
+            value = annotate_with_constraint(value, constraint)
+            return Composite(value, varname, node)
         else:
             # not sure when (if ever) the other contexts can happen
             self.show_error(node, f"Bad context: {node.ctx}", ErrorCode.unexpected_node)
@@ -4056,6 +4082,10 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor, CanAssignContext):
         elif isinstance(node, (ast.ExtSlice, ast.Slice)):
             # These don't have a .lineno attribute, which would otherwise cause trouble.
             composite = Composite(self.visit(node), None, None)
+        # We need better support for version-straddling code
+        # static analysis: ignore[value_always_true]
+        elif hasattr(ast, "NamedExpr") and isinstance(node, ast.NamedExpr):
+            composite = self.composite_from_walrus(node)
         else:
             composite = Composite(self.visit(node), None, node)
         if self.annotate:
diff --git a/pyanalyze/stacked_scopes.py b/pyanalyze/stacked_scopes.py
index c01b2c0..e44ee55 100644
--- a/pyanalyze/stacked_scopes.py
+++ b/pyanalyze/stacked_scopes.py
@@ -721,7 +721,7 @@ class Scope:
 
     def set(
         self, varname: Varname, value: Value, node: Node, state: VisitorState
-    ) -> None:
+    ) -> VarnameOrigin:
         if varname not in self:
             self.variables[varname] = value
         elif isinstance(value, AnyValue) or not safe_equals(
@@ -741,6 +741,7 @@ class Scope:
                 self.variables[varname] = value
             else:
                 self.variables[varname] = unite_values(existing, value)
+        return EMPTY_ORIGIN
 
     def items(self) -> Iterable[Tuple[Varname, Value]]:
         return self.variables.items()
@@ -1027,10 +1028,10 @@ class FunctionScope(Scope):
 
     def set(
         self, varname: Varname, value: Value, node: Node, state: VisitorState
-    ) -> None:
+    ) -> VarnameOrigin:
         if isinstance(value, ReferencingValue):
             self.referencing_value_vars[varname] = value
-            return
+            return EMPTY_ORIGIN
         ref_var = self.referencing_value_vars[varname]
         if isinstance(ref_var, ReferencingValue):
             ref_var.scope.set(ref_var.name, value, node, state)
@@ -1047,6 +1048,7 @@ class FunctionScope(Scope):
             self.name_to_current_definition_nodes[composite] = []
         self.name_to_all_definition_nodes[varname].add(node)
         self._add_composite(varname)
+        return frozenset([node])
 
     def get_local(
         self,


Test patch:
diff --git a/pyanalyze/test_ast_annotator.py b/pyanalyze/test_ast_annotator.py
index 4a162c1..071a46d 100644
--- a/pyanalyze/test_ast_annotator.py
+++ b/pyanalyze/test_ast_annotator.py
@@ -4,7 +4,7 @@ from typing import Callable, Type
 
 from .ast_annotator import annotate_file, annotate_code
 from .analysis_lib import files_with_extension_from_directory
-from .value import KnownValue, Value
+from .value import KnownValue, Value, unannotate
 
 
 def _check_inferred_value(
@@ -16,7 +16,7 @@ def _check_inferred_value(
     for node in ast.walk(tree):
         if isinstance(node, node_type) and predicate(node):
             assert hasattr(node, "inferred_value"), repr(node)
-            assert value == node.inferred_value, ast.dump(node)
+            assert value == unannotate(node.inferred_value), ast.dump(node)
 
 
 def test_annotate_code() -> None:
@@ -26,37 +26,37 @@ def test_annotate_code() -> None:
 
     tree = annotate_code(
         """
-class X:
-    def __init__(self):
-        self.a = 1
-    """
+        class X:
+            def __init__(self):
+                self.a = 1
+        """
     )
     _check_inferred_value(tree, ast.Attribute, KnownValue(1))
     tree = annotate_code(
         """
-class X:
-    def __init__(self):
-        self.a = 1
+        class X:
+            def __init__(self):
+                self.a = 1
 
-x = X()
-x.a + 1
-    """
+        x = X()
+        x.a + 1
+        """
     )
     _check_inferred_value(tree, ast.BinOp, KnownValue(2))
 
     tree = annotate_code(
         """
-class A:
-    def __init__(self):
-        self.a = 1
+        class A:
+            def __init__(self):
+                self.a = 1
 
-    def bla(self):
-        return self.a
+            def bla(self):
+                return self.a
 
 
-a = A()
-b = a.bla()
-"""
+        a = A()
+        b = a.bla()
+        """
     )
     _check_inferred_value(tree, ast.Name, KnownValue(1), lambda node: node.id == "b")
 
diff --git a/pyanalyze/test_name_check_visitor.py b/pyanalyze/test_name_check_visitor.py
index 8f1a492..7891810 100644
--- a/pyanalyze/test_name_check_visitor.py
+++ b/pyanalyze/test_name_check_visitor.py
@@ -43,7 +43,7 @@ from .value import (
 )
 from .tests import proxied_fn, autogenerated, l0cached_async_fn, PropertyObject
 from . import test_node_visitor
-from .test_node_visitor import assert_passes, assert_fails, only_before
+from .test_node_visitor import assert_passes, assert_fails, only_before, skip_before
 
 
 # ===================================================
@@ -2577,3 +2577,38 @@ class TestIncompatibleOverride(TestNameCheckVisitorBase):
 
             def pacarana(self, b: int) -> None:
                 pass
+
+
+class TestWalrus(TestNameCheckVisitorBase):
+    @skip_before((3, 8))
+    def test(self):
+        self.assert_passes(
+            """
+            from typing import Optional
+
+            def opt() -> Optional[int]:
+                return None
+
+            def capybara():
+                if (x := opt()):
+                    assert_is_value(x, TypedValue(int))
+                assert_is_value(x, TypedValue(int) | KnownValue(None))
+
+                if (y := opt()) is not None:
+                    assert_is_value(y, TypedValue(int))
+                assert_is_value(y, TypedValue(int) | KnownValue(None))
+            """
+        )
+
+    @skip_before((3, 8))
+    def test_comprehension_scope(self):
+        self.assert_passes(
+            """
+            from typing import List, Optional
+
+            def capybara(elts: List[Optional[int]]) -> None:
+                if any((x := i) is not None for i in elts):
+                    assert_is_value(x, TypedValue(int) | KnownValue(None))
+                    print(i)  # E: undefined_name
+            """
+        )