Improve handling of forward references (#364)

As discussed [here](#276)

Fixes #276
Fixes #317
Fixes #278

Co-authored-by: Alex Waygood <[email protected]>
Co-authored-by: Akuli <[email protected]>

CHANGELOG.md

@@ -3,6 +3,27 @@
 ## Unreleased
 
 * flake8-pyi no longer supports being run with flake8 <5.0.4.
+* The way in which flake8-pyi modifies pyflakes runs has been improved:
+  * When flake8-pyi is installed, pyflakes now correctly recognises an annotation as
+    being equivalent to a binding assignment in a stub file, reducing false
+    positives from flake8's F821 error code.
+  * When flake8-pyi is installed, there are now fewer pyflakes positives from class
+    definitions that have forward references in the bases tuple for the purpose of
+    creating recursive or circular type definitions. These are invalid in `.py` files,
+    but are supported in stub files.
+  * When flake8-pyi is installed, pyflakes will also *complain* about code which (in
+    combination with flake8-pyi) it previously had no issue with. For example, it will
+    now complain about this code:
+
+    ```py
+    class Foo(Bar): ...
+    class Bar: ...
+    ```
+
+    Although the above code is legal in a stub file, it is considered poor style, and
+    the forward reference serves no purpose (there is no recursive or circular
+    definition). As such, it is now disallowed by pyflakes when flake8-pyi is
+    installed.
 
 ## 23.4.1
 

pyi.py

@@ -21,7 +21,7 @@
 from flake8 import checker  # type: ignore[import]
 from flake8.options.manager import OptionManager  # type: ignore[import]
 from flake8.plugins.pyflakes import FlakesChecker  # type: ignore[import]
-from pyflakes.checker import ClassDefinition, ClassScope, FunctionScope, ModuleScope
+from pyflakes.checker import ModuleScope
 
 if sys.version_info >= (3, 9):
     from ast import unparse
@@ -159,7 +159,41 @@ class TypeVarInfo(NamedTuple):
 )
 
 
+class PyflakesPreProcessor(ast.NodeTransformer):
+    """Transform AST prior to passing it to pyflakes.
+
+    This reduces false positives on recursive class definitions.
+    """
+
+    def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
+        self.generic_visit(node)
+        node.bases = [
+            # Remove the subscript to prevent F821 errors from being raised
+            # for (valid) recursive definitions: Foo[Bar] --> Foo
+            base.value if isinstance(base, ast.Subscript) else base
+            for base in node.bases
+        ]
+        return node
+
+
 class PyiAwareFlakesChecker(FlakesChecker):
+    def __init__(self, tree: ast.AST, *args: Any, **kwargs: Any) -> None:
+        super().__init__(PyflakesPreProcessor().visit(tree), *args, **kwargs)
+
+    @property
+    def annotationsFutureEnabled(self):
+        """pyflakes can already handle forward refs for annotations, but only via
+        `from __future__ import annotations`.
+
+        We don't want to bother including this in every file, so we just set this to `True`.
+        """
+        return True
+
+    @annotationsFutureEnabled.setter
+    def annotationsFutureEnabled(self, value: bool):
+        """Does nothing, as we always want this property to be `True`."""
+        pass
+
     def deferHandleNode(self, node: ast.AST | None, parent) -> None:
         self.deferFunction(lambda: self.handleNode(node, parent))
 
@@ -182,29 +216,6 @@ def ASSIGN(
 
         self.deferHandleNode(tree.value, tree)
 
-    def ANNASSIGN(self, node: ast.AnnAssign) -> None:
-        """
-        Annotated assignments don't have annotations evaluated on function
-        scope, hence the custom implementation. Compared to the pyflakes
-        version, we defer evaluation of the annotations (and values on
-        module level).
-        """
-        if node.value:
-            # Only bind the *target* if the assignment has value.
-            # Otherwise it's not really ast.Store and shouldn't silence
-            # UndefinedLocal warnings.
-            self.handleNode(node.target, node)
-        if not isinstance(self.scope, FunctionScope):
-            self.deferHandleNode(node.annotation, node)
-        if node.value:
-            # If the assignment has value, handle the *value*...
-            if isinstance(self.scope, ModuleScope):
-                # ...later (if module scope).
-                self.deferHandleNode(node.value, node)
-            else:
-                # ...now.
-                self.handleNode(node.value, node)
-
     def LAMBDA(self, node: ast.Lambda) -> None:
         """This is likely very brittle, currently works for pyflakes 1.3.0.
 
@@ -216,35 +227,6 @@ def LAMBDA(self, node: ast.Lambda) -> None:
         super().LAMBDA(node)
         self.handleNode, self.deferHandleNode = self.deferHandleNode, self.handleNode  # type: ignore[method-assign]
 
-    def CLASSDEF(self, node: ast.ClassDef) -> None:
-        if not isinstance(self.scope, ModuleScope):
-            # This shouldn't be necessary because .pyi files don't nest
-            # scopes much, but better safe than sorry.
-            super().CLASSDEF(node)
-            return
-
-        # What follows is copied from pyflakes 1.3.0. The only changes are the
-        # deferHandleNode calls.
-        for decorator in node.decorator_list:
-            self.handleNode(decorator, node)
-        for baseNode in node.bases:
-            self.deferHandleNode(baseNode, node)
-        for keywordNode in node.keywords:
-            self.deferHandleNode(keywordNode, node)
-        self.pushScope(ClassScope)
-        # doctest does not process doctest within a doctest
-        # classes within classes are processed.
-        if (
-            self.withDoctest
-            and not self._in_doctest()
-            and not isinstance(self.scope, FunctionScope)
-        ):
-            self.deferFunction(lambda: self.handleDoctests(node))
-        for stmt in node.body:
-            self.handleNode(stmt, node)
-        self.popScope()
-        self.addBinding(node, ClassDefinition(node.name, node))
-
     def handleNodeDelete(self, node: ast.AST) -> None:
         """Null implementation.
 

tests/forward_refs.pyi

@@ -8,11 +8,16 @@ __author__: str = ...
 
 def make_default_c() -> C: ...
 
-class D(C):
+# Disallow forward refs for base classes
+class D(C):  # F821 undefined name 'C'
     parent: C
     def __init__(self) -> None: ...
 
 class C:
     other: C
     def __init__(self) -> None: ...
     def from_str(self, s: str) -> C: ...
+
+class Outer:
+    class Inner1: ...
+    class Inner2(list[Outer.Inner1]): ...

tests/forward_refs_annassign.pyi

@@ -8,11 +8,49 @@ __author__: str
 
 def make_default_c() -> C: ...
 
-class D(C):
+# Disallow forward refs for base classes
+class D(C):  # F821 undefined name 'C'
     parent: C
     def __init__(self) -> None: ...
 
 class C:
     other: C = ...
     def __init__(self) -> None: ...
     def from_str(self, s: str) -> C: ...
+
+class Baz(metaclass=Meta):  # F821 undefined name 'Meta'
+    ...
+
+class Foo(Bar, Baz, metaclass=Meta):  # F821 undefined name 'Bar'  # F821 undefined name 'Meta'
+    ...
+
+class Meta(type):
+    ...
+
+class Bar(metaclass=Meta):
+    ...
+
+# Allow circular references in annotations
+class A:
+    foo: B
+    bar: dict[str, B]
+
+class B:
+    foo: A
+    bar: dict[str, A]
+
+# This recursive definition is allowed.
+# We allow it by disabling checks for anything in the subscript
+class Leaf: ...
+class Tree(list[Tree | Leaf]): ...
+
+# This recursive definition is not allowed.
+class Parent(Child): ...   # F821 undefined name 'Child'
+class Child(Parent): ...
+
+class MyClass:
+    foo: int
+    bar = foo
+
+baz: MyClass
+eggs = baz