Repository: quora/pyanalyze
Pull request #614 (2023-05-02 01:23:46)

Add support for PEP 698 @override


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index b1013ad..b647018 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -13,6 +13,7 @@
 - Fix bug where objects with a `__call__` method that takes `*args` instead
   of `self` were not considered callable (#600)
 - Better typechecking support for async generators (#594)
+- Add support for the PEP 698 `@override` decorator (#614)
 
 ## Version 0.9.0 (January 16, 2023)
 
diff --git a/pyanalyze/error_code.py b/pyanalyze/error_code.py
index 337af42..61285ea 100644
--- a/pyanalyze/error_code.py
+++ b/pyanalyze/error_code.py
@@ -102,6 +102,8 @@ class ErrorCode(enum.Enum):
     invalid_import = 82
     too_many_positional_args = 83
     deprecated = 84
+    invalid_override_decorator = 85
+    override_does_not_override = 86
 
 
 # Allow testing unannotated functions without too much fuss
@@ -228,6 +230,8 @@ ERROR_DESCRIPTION = {
     ErrorCode.too_many_positional_args: (
         "Call with many positional arguments should use keyword arguments"
     ),
+    ErrorCode.invalid_override_decorator: "@override decorator in invalid location",
+    ErrorCode.override_does_not_override: "Method does not override any base method",
 }
 
 
diff --git a/pyanalyze/functions.py b/pyanalyze/functions.py
index b5a74c9..de2ad95 100644
--- a/pyanalyze/functions.py
+++ b/pyanalyze/functions.py
@@ -21,6 +21,7 @@ from .error_code import ErrorCode
 from .extensions import evaluated, overload, real_overload
 from .node_visitor import ErrorContext
 from .options import Options, PyObjectSequenceOption
+from .safe import is_typing_name
 from .signature import ParameterKind, Signature, SigParameter
 from .stacked_scopes import Composite
 from .value import (
@@ -82,6 +83,7 @@ class FunctionInfo:
     is_staticmethod: bool  # has @staticmethod
     is_decorated_coroutine: bool  # has @asyncio.coroutine
     is_overload: bool  # typing.overload or pyanalyze.extensions.overload
+    is_override: bool  # @typing.override
     is_evaluated: bool  # @pyanalyze.extensions.evaluated
     is_abstractmethod: bool  # has @abstractmethod
     # a list of tuples of (decorator function, applied decorator function, AST node). These are
@@ -219,6 +221,7 @@ def compute_function_info(
     is_decorated_coroutine = False
     is_staticmethod = False
     is_overload = False
+    is_override = False
     is_abstractmethod = False
     is_evaluated = False
     decorators = []
@@ -242,22 +245,24 @@ def compute_function_info(
             decorators.append((callee, decorator_value, decorator))
         else:
             decorator_value = ctx.visit_expression(decorator)
-            if decorator_value == KnownValue(classmethod):
-                is_classmethod = True
-            elif decorator_value == KnownValue(staticmethod):
-                is_staticmethod = True
-            elif sys.version_info < (3, 11) and decorator_value == KnownValue(
-                asyncio.coroutine  # static analysis: ignore[undefined_attribute]
-            ):
-                is_decorated_coroutine = True
-            elif decorator_value == KnownValue(
-                real_overload
-            ) or decorator_value == KnownValue(overload):
-                is_overload = True
-            elif decorator_value == KnownValue(abstractmethod):
-                is_abstractmethod = True
-            elif decorator_value == KnownValue(evaluated):
-                is_evaluated = True
+            if isinstance(decorator_value, KnownValue):
+                val = decorator_value.val
+                if val is classmethod:
+                    is_classmethod = True
+                elif val is staticmethod:
+                    is_staticmethod = True
+                elif sys.version_info < (3, 11) and val is (
+                    asyncio.coroutine  # static analysis: ignore[undefined_attribute]
+                ):
+                    is_decorated_coroutine = True
+                elif val is real_overload or val is overload:
+                    is_overload = True
+                elif val is abstractmethod:
+                    is_abstractmethod = True
+                elif val is evaluated:
+                    is_evaluated = True
+                elif is_typing_name(val, "override"):
+                    is_override = True
             decorators.append((decorator_value, decorator_value, decorator))
     params = compute_parameters(
         node,
@@ -278,6 +283,7 @@ def compute_function_info(
         is_staticmethod=is_staticmethod,
         is_abstractmethod=is_abstractmethod,
         is_overload=is_overload,
+        is_override=is_override,
         is_evaluated=is_evaluated,
         decorators=decorators,
         node=node,
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index e9dcdf7..5995f36 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -1459,13 +1459,11 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         origin = current_scope.set(varname, value, lookup_node, self.state)
         return value, origin
 
-    def _check_for_incompatible_overrides(
-        self, varname: str, node: ast.AST, value: Value
-    ) -> None:
+    def _get_base_class_attributes(
+        self, varname: str, node: ast.AST
+    ) -> Iterable[Tuple[Union[type, str], Value]]:
         if self.current_class is None:
             return
-        if varname in self.options.get_value_for(IgnoredForIncompatibleOverride):
-            return
         for base_class in self.get_generic_bases(self.current_class):
             if base_class is self.current_class:
                 continue
@@ -1480,6 +1478,17 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                 record_reads=False,
             )
             base_value = attributes.get_attribute(ctx)
+            if base_value is not UNINITIALIZED_VALUE:
+                yield base_class, base_value
+
+    def _check_for_incompatible_overrides(
+        self, varname: str, node: ast.AST, value: Value
+    ) -> None:
+        if self.current_class is None:
+            return
+        if varname in self.options.get_value_for(IgnoredForIncompatibleOverride):
+            return
+        for base_class, base_value in self._get_base_class_attributes(varname, node):
             can_assign = self._can_assign_to_base(base_value, value)
             if isinstance(can_assign, CanAssignError):
                 error = CanAssignError(
@@ -1773,6 +1782,17 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
             else:
                 self._show_error_if_checking(node, error_code=ErrorCode.missing_return)
 
+        if info.is_override:
+            if self.current_class is None:
+                self._show_error_if_checking(
+                    node, error_code=ErrorCode.invalid_override_decorator
+                )
+            else:
+                if not any(self._get_base_class_attributes(node.name, node)):
+                    self._show_error_if_checking(
+                        node, error_code=ErrorCode.override_does_not_override
+                    )
+
         if node.returns is None:
             if (
                 result.has_return


Test patch:
diff --git a/pyanalyze/test_override.py b/pyanalyze/test_override.py
new file mode 100644
index 0000000..6b24e8c
--- /dev/null
+++ b/pyanalyze/test_override.py
@@ -0,0 +1,40 @@
+# static analysis: ignore
+from .error_code import ErrorCode
+from .test_name_check_visitor import TestNameCheckVisitorBase
+from .test_node_visitor import assert_fails, assert_passes
+
+
+class TestOverride(TestNameCheckVisitorBase):
+    @assert_fails(ErrorCode.invalid_override_decorator)
+    def test_invalid_usage(self):
+        from typing_extensions import override
+
+        @override
+        def not_a_method():
+            pass
+
+    @assert_passes()
+    def test_valid_method(self):
+        from typing_extensions import override
+
+        class Base:
+            def method(self):
+                pass
+
+        class Capybara(Base):
+            @override
+            def method(self):
+                pass
+
+    @assert_fails(ErrorCode.override_does_not_override)
+    def test_invalid_method(self):
+        from typing_extensions import override
+
+        class Base:
+            def method(self):
+                pass
+
+        class Capybara(Base):
+            @override
+            def no_base_method(self):  # E: override_does_not_override
+                pass