Repository: quora/pyanalyze
Pull request #505 (2022-03-20 01:54:01)

Support declared types


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 58b9c74..0277a8c 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Basic support for checking `Final` and for checking re-assignments
+  to variables declared with a specific type (#505)
 - Correctly check the `self` argument to `@property` getters (#506)
 - Correctly track assignments of variables inside `try` blocks
   and inside `with` blocks that may suppress exceptions (#504)
diff --git a/pyanalyze/error_code.py b/pyanalyze/error_code.py
index d74150f..1699a7a 100644
--- a/pyanalyze/error_code.py
+++ b/pyanalyze/error_code.py
@@ -93,6 +93,8 @@ class ErrorCode(enum.Enum):
     bad_match = 75
     bad_evaluator = 76
     implicit_any = 77
+    already_declared = 78
+    invalid_annotated_assignment = 79
 
 
 # Allow testing unannotated functions without too much fuss
@@ -212,6 +214,8 @@ ERROR_DESCRIPTION = {
     ErrorCode.bad_match: "Invalid type in match statement",
     ErrorCode.bad_evaluator: "Invalid code in type evaluator",
     ErrorCode.implicit_any: "Value is inferred as Any",
+    ErrorCode.already_declared: "Name is already declared",
+    ErrorCode.invalid_annotated_assignment: "Invalid annotated assignment",
 }
 
 
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index 37cdcf7..14a05a1 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -996,7 +996,8 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
     options: Options
     arg_spec_cache: ArgSpecCache
     reexport_tracker: ImplicitReexportTracker
-    being_assigned: Value
+    being_assigned: Optional[Value]
+    ann_assign_type: Optional[Tuple[Optional[Value], bool]]
     current_class: Optional[type]
     current_function_name: Optional[str]
     current_enum_members: Optional[Dict[object, str]]
@@ -1037,6 +1038,7 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         self.state = VisitorState.collect_names
         # value currently being assigned
         self.being_assigned = AnyValue(AnySource.inference)
+        self.ann_assign_type = None
         # current match target
         self.match_subject = Composite(AnyValue(AnySource.inference))
         # current class (for inferring the type of cls and self arguments)
@@ -1335,12 +1337,44 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
     def _set_name_in_scope(
         self,
         varname: str,
-        node: object,
+        node: ast.AST,
         value: Value = AnyValue(AnySource.inference),
         *,
         private: bool = False,
+        lookup_node: object = None,
     ) -> Tuple[Value, VarnameOrigin]:
+        if lookup_node is None:
+            lookup_node = node
+
         current_scope = self.scopes.current_scope()
+        if self.ann_assign_type is not None:
+            expected_type, is_final = self.ann_assign_type
+            if not current_scope.set_declared_type(
+                varname, expected_type, is_final, node
+            ):
+                self._show_error_if_checking(
+                    node, f"{varname} already declared", ErrorCode.already_declared
+                )
+
+        else:
+            declared_type = current_scope.get_declared_type(varname)
+            if declared_type is not None:
+                can_assign = declared_type.can_assign(value, self)
+                if isinstance(can_assign, CanAssignError):
+                    self._show_error_if_checking(
+                        node,
+                        f"Incompatible assignment: expected {declared_type}, got"
+                        f" {value}",
+                        error_code=ErrorCode.incompatible_assignment,
+                        detail=can_assign.display(),
+                    )
+            if current_scope.is_final(varname):
+                self._show_error_if_checking(
+                    node,
+                    f"Cannot assign to final name {varname}",
+                    ErrorCode.incompatible_assignment,
+                )
+
         scope_type = current_scope.scope_type
         if self.module is not None and scope_type == ScopeType.module_scope:
             if self.module.__name__ is not None and not private:
@@ -1348,12 +1382,12 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                     self.module.__name__, varname
                 )
             if varname in current_scope:
-                value, _ = current_scope.get_local(varname, node, self.state)
+                value, _ = current_scope.get_local(varname, lookup_node, self.state)
                 return value, EMPTY_ORIGIN
-        if scope_type == ScopeType.class_scope and isinstance(node, ast.AST):
+        if scope_type == ScopeType.class_scope:
             self._check_for_incompatible_overrides(varname, node, value)
             self._check_for_class_variable_redefinition(varname, node)
-        origin = current_scope.set(varname, value, node, self.state)
+        origin = current_scope.set(varname, value, lookup_node, self.state)
         return value, origin
 
     def _check_for_incompatible_overrides(
@@ -1418,7 +1452,7 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
     def _check_for_class_variable_redefinition(
         self, varname: str, node: ast.AST
     ) -> None:
-        if varname not in self.scopes.current_scope():
+        if varname not in self.scopes.current_scope().variables:
             return
 
         # exclude cases where we do @<property>.setter
@@ -2273,7 +2307,11 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
             ):
                 continue
             self._set_name_in_scope(
-                name, (node, name), KnownValue(value), private=not force_public
+                name,
+                node,
+                KnownValue(value),
+                private=not force_public,
+                lookup_node=(node, name),
             )
 
     def _handle_imports(
@@ -2627,6 +2665,9 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                 else:
                     target_length += 1
 
+            assert (
+                self.being_assigned is not None
+            ), "annotated assignment can only have a single target"
             being_assigned = unpack_values(
                 self.being_assigned, self, target_length, post_starred_length
             )
@@ -3756,40 +3797,45 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         if isinstance(annotation, KnownValue) and is_typing_name(
             annotation.val, "Final"
         ):
-            # TODO disallow assignments to Final variables (current code
-            # just avoids false positive errors).
             is_final = True
-            expected_type = AnyValue(AnySource.marker)
+            expected_type = None
         else:
             expected_type = self._value_of_annotation_type(
                 annotation,
                 node.annotation,
                 is_typeddict=self.is_in_typeddict_definition(),
             )
+            # TODO: Also extract Final from more complex annotations
             is_final = False
 
-        if node.value:
+        if node.value is not None:
             is_yield = isinstance(node.value, ast.Yield)
             value = self.visit(node.value)
-            can_assign = expected_type.can_assign(value, self)
-            if isinstance(can_assign, CanAssignError):
-                self._show_error_if_checking(
-                    node.value,
-                    f"Incompatible assignment: expected {expected_type}, got {value}",
-                    error_code=ErrorCode.incompatible_assignment,
-                    detail=can_assign.display(),
-                )
 
-            with qcore.override(
-                self, "being_assigned", value if is_final else expected_type
-            ), self.yield_checker.check_yield_result_assignment(is_yield):
-                self.visit(node.target)
+            if expected_type is not None:
+                can_assign = expected_type.can_assign(value, self)
+                if isinstance(can_assign, CanAssignError):
+                    self._show_error_if_checking(
+                        node,
+                        f"Incompatible assignment: expected {expected_type}, got"
+                        f" {value}",
+                        error_code=ErrorCode.incompatible_assignment,
+                        detail=can_assign.display(),
+                    )
+                # We set the declared type on initial assignment, so that the
+                # annotation can be used to adjust pyanalyze's type inference.
+                value = expected_type
+
         else:
-            with qcore.override(self, "being_assigned", expected_type):
-                self.visit(node.target)
-        # TODO: Idea for what to do if there is no value:
-        # - Scopes keep track of a map {name: expected type}
-        # - Assignments that are inconsistent with the declared type produce an error.
+            is_yield = False
+            value = None
+
+        with qcore.override(
+            self, "being_assigned", value
+        ), self.yield_checker.check_yield_result_assignment(is_yield), qcore.override(
+            self, "ann_assign_type", (expected_type, is_final)
+        ):
+            self.visit(node.target)
 
     def visit_AugAssign(self, node: ast.AugAssign) -> None:
         is_yield = isinstance(node.value, ast.Yield)
@@ -3829,18 +3875,22 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                     (ast.stmt, ast.comprehension)
                 )
                 self._name_node_to_statement[node] = statement
-                # If we're in an AnnAssign without a value, we skip the assignment,
-                # since no value is actually assigned to the name.
-                is_ann_assign = (
-                    isinstance(statement, ast.AnnAssign) and statement.value is None
-                )
-            else:
-                is_ann_assign = False
+
             value = self.being_assigned
-            origin = EMPTY_ORIGIN
-            if not is_ann_assign:
+            if (
+                value is None
+                and self.scopes.current_scope().scope_type == ScopeType.class_scope
+                and self.ann_assign_type is not None
+            ):
+                ann_assign_type, _ = self.ann_assign_type
+                if ann_assign_type is not None:
+                    value = ann_assign_type
+            if value is not None:
                 self.yield_checker.record_assignment(node.id)
                 value, origin = self._set_name_in_scope(node.id, node, value=value)
+            else:
+                value = AnyValue(AnySource.inference)
+                origin = EMPTY_ORIGIN
             varname = VarnameWithOrigin(node.id, origin)
             constraint = Constraint(varname, ConstraintType.is_truthy, True, None)
             value = annotate_with_constraint(value, constraint)
@@ -3856,6 +3906,8 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
 
     def visit_arg(self, node: ast.arg) -> None:
         self.yield_checker.record_assignment(node.arg)
+        # it's none only for AnnAssign nodes without a value
+        assert self.being_assigned is not None, "should not happen"
         self._set_name_in_scope(node.arg, node, value=self.being_assigned)
 
     def _should_use_varname_value(self, value: Value) -> bool:
@@ -3913,6 +3965,15 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         index = index_composite.value
 
         if isinstance(node.ctx, ast.Store):
+            if self.ann_assign_type is not None:
+                self._show_error_if_checking(
+                    node, error_code=ErrorCode.invalid_annotated_assignment
+                )
+            if self.being_assigned is None:
+                assert (
+                    self.ann_assign_type is not None
+                ), "should only happen in AnnAssign"
+                return AnyValue(AnySource.inference)
             if (
                 composite_var is not None
                 and self.scopes.scope_type() == ScopeType.function_scope
@@ -4108,6 +4169,13 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         root_composite = self.composite_from_node(node.value)
         composite = self._extend_composite(root_composite, node.attr, node)
         if self._is_write_ctx(node.ctx):
+            # TODO: We should do something here if we're in an AnnAssign, e.g.
+            # note the type in the class's namespace.
+            if self.being_assigned is None:
+                assert (
+                    self.ann_assign_type is not None
+                ), "should only happen in AnnAssign"
+                return Composite(AnyValue(AnySource.inference), composite, node)
             if (
                 composite is not None
                 and self.scopes.scope_type() == ScopeType.function_scope
@@ -4513,7 +4581,7 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                         )
                     return_value = KnownValue(result)
 
-        if return_value is NO_RETURN_VALUE:
+        if return_value is NO_RETURN_VALUE and node is not None:
             self._set_name_in_scope(LEAVES_SCOPE, node, AnyValue(AnySource.marker))
 
         # for .asynq functions, we use the argspec for the underlying function, but that means
diff --git a/pyanalyze/stacked_scopes.py b/pyanalyze/stacked_scopes.py
index bd39239..4cccfbf 100644
--- a/pyanalyze/stacked_scopes.py
+++ b/pyanalyze/stacked_scopes.py
@@ -709,6 +709,9 @@ class Scope:
     scope_node: Optional[Node] = None
     scope_object: Optional[object] = None
     simplification_limit: Optional[int] = None
+    declared_types: Dict[str, Tuple[Optional[Value], bool, AST]] = field(
+        default_factory=dict
+    )
 
     def __post_init__(self) -> None:
         if self.parent_scope is not None:
@@ -767,7 +770,7 @@ class Scope:
     def set(
         self, varname: Varname, value: Value, node: Node, state: VisitorState
     ) -> VarnameOrigin:
-        if varname not in self:
+        if varname not in self.variables:
             self.variables[varname] = value
         elif isinstance(value, AnyValue) or not safe_equals(
             self.variables[varname], value
@@ -794,8 +797,35 @@ class Scope:
     def all_variables(self) -> Iterable[Varname]:
         return self.variables
 
+    def set_declared_type(
+        self, varname: str, typ: Optional[Value], is_final: bool, node: AST
+    ) -> bool:
+        if varname in self.declared_types:
+            _, _, existing_node = self.declared_types[varname]
+            already_present = node is not existing_node
+            # Don't replace the existing node, or we'll generate spurious already_declared
+            # errors.
+            node = existing_node
+        else:
+            already_present = False
+        # Even if we give an error, still honor the later type.
+        self.declared_types[varname] = (typ, is_final, node)
+        return not already_present
+
+    def get_declared_type(self, varname: str) -> Optional[Value]:
+        if varname not in self.declared_types:
+            return None
+        typ, _, _ = self.declared_types[varname]
+        return typ
+
+    def is_final(self, varname: str) -> bool:
+        if varname not in self.declared_types:
+            return False
+        _, is_final, _ = self.declared_types[varname]
+        return is_final
+
     def __contains__(self, varname: Varname) -> bool:
-        return varname in self.variables
+        return varname in self.variables or varname in self.declared_types
 
     @contextlib.contextmanager
     def suppressing_subscope(self) -> Iterator[SubScope]:
@@ -1338,7 +1368,10 @@ class FunctionScope(Scope):
         yield from self.name_to_current_definition_nodes
 
     def __contains__(self, varname: Varname) -> bool:
-        return varname in self.name_to_all_definition_nodes
+        return (
+            varname in self.name_to_all_definition_nodes
+            or varname in self.declared_types
+        )
 
 
 class StackedScopes:


Test patch:
diff --git a/pyanalyze/test_implementation.py b/pyanalyze/test_implementation.py
index 6b48252..0a6a7eb 100644
--- a/pyanalyze/test_implementation.py
+++ b/pyanalyze/test_implementation.py
@@ -395,15 +395,15 @@ class TestGenericMutators(TestNameCheckVisitorBase):
                 lst, SequenceIncompleteValue(list, [TypedValue(int), KnownValue(1)])
             )
 
-            lst: List[str] = ["x"]
-            assert_is_value(lst, GenericValue(list, [TypedValue(str)]))
-            lst.append("y")
-            assert_is_value(lst, GenericValue(list, [TypedValue(str)]))
+            lst2: List[str] = ["x"]
+            assert_is_value(lst2, GenericValue(list, [TypedValue(str)]))
+            lst2.append("y")
+            assert_is_value(lst2, GenericValue(list, [TypedValue(str)]))
 
-            lst = ["x"]
-            assert_is_value(lst, KnownValue(["x"]))
-            lst.append(3)
-            assert_is_value(lst, KnownValue(["x", 3]))
+            lst3 = ["x"]
+            assert_is_value(lst3, KnownValue(["x"]))
+            lst3.append(3)
+            assert_is_value(lst3, KnownValue(["x", 3]))
 
     @assert_passes()
     def test_list_append_pos_only(self):
@@ -426,17 +426,17 @@ class TestGenericMutators(TestNameCheckVisitorBase):
         from typing import Set
 
         def capybara(x: int):
-            lst = {x}
-            assert_is_value(lst, SequenceIncompleteValue(set, [TypedValue(int)]))
-            lst.add(1)
+            s = {x}
+            assert_is_value(s, SequenceIncompleteValue(set, [TypedValue(int)]))
+            s.add(1)
             assert_is_value(
-                lst, SequenceIncompleteValue(set, [TypedValue(int), KnownValue(1)])
+                s, SequenceIncompleteValue(set, [TypedValue(int), KnownValue(1)])
             )
 
-            lst: Set[str] = {"x"}
-            assert_is_value(lst, GenericValue(set, [TypedValue(str)]))
-            lst.add("y")
-            assert_is_value(lst, GenericValue(set, [TypedValue(str)]))
+            s2: Set[str] = {"x"}
+            assert_is_value(s2, GenericValue(set, [TypedValue(str)]))
+            s2.add("y")
+            assert_is_value(s2, GenericValue(set, [TypedValue(str)]))
 
     @assert_passes()
     def test_list_add(self):
@@ -546,10 +546,10 @@ class TestGenericMutators(TestNameCheckVisitorBase):
                 ),
             )
 
-            lst: List[int] = [3]
-            assert_is_value(lst, GenericValue(list, [TypedValue(int)]))
-            lst += [x]
-            assert_is_value(lst, GenericValue(list, [TypedValue(int)]))
+            lst2: List[int] = [3]
+            assert_is_value(lst2, GenericValue(list, [TypedValue(int)]))
+            lst2 += [x]
+            assert_is_value(lst2, GenericValue(list, [TypedValue(int)]))
 
     @assert_passes()
     def test_list_iadd_never(self):
diff --git a/pyanalyze/test_name_check_visitor.py b/pyanalyze/test_name_check_visitor.py
index e1c4f87..6f83f32 100644
--- a/pyanalyze/test_name_check_visitor.py
+++ b/pyanalyze/test_name_check_visitor.py
@@ -1836,7 +1836,7 @@ class TestRequireAnnotations(TestNameCheckVisitorBase):
 
 class TestAnnAssign(TestNameCheckVisitorBase):
     @assert_passes()
-    def test(self):
+    def test_simple(self):
         from typing_extensions import Final
 
         def capybara() -> None:
@@ -1847,6 +1847,51 @@ class TestAnnAssign(TestNameCheckVisitorBase):
             z: float
             print(z)  # E: undefined_name
 
+            y: float = 4.0  # E: already_declared
+            assert_is_value(y, TypedValue(float))
+
+    @assert_passes()
+    def test_final(self):
+        from typing_extensions import Final
+
+        x: Final = 1
+        x = 2  # E: incompatible_assignment
+
+        def capybara():
+            y: Final = 1
+            y = 2  # E: incompatible_assignment
+            return y
+
+    @assert_passes()
+    def test_inconsistent_type(self):
+        def capybara():
+            x: int = 1
+            assert_is_value(x, TypedValue(int))
+            x = "x"  # E: incompatible_assignment
+
+            y: int = "y"  # E: incompatible_assignment
+            return (x, y)
+
+    @assert_passes()
+    def test_class_scope(self):
+        class Capybara:
+            x: int = 0
+            assert_is_value(x, TypedValue(int))
+
+            def __init__(self) -> None:
+                self.y: object = 3
+
+            def method(self):
+                assert_is_value(self.y, TypedValue(object))
+
+    @assert_passes()
+    def test_loop(self):
+        def capybara():
+            for i in range(3):
+                j: int = i
+
+            j: int = 0  # E: already_declared
+
 
 class TestWhile(TestNameCheckVisitorBase):
     @assert_passes()