fix: Make sure to visit returns not just the argument annotations

snok · Dec 13, 2024 · eef57df · eef57df
1 parent 40c43b5
commit eef57df
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 19 deletions.
diff --git a/flake8_type_checking/checker.py b/flake8_type_checking/checker.py
@@ -291,6 +291,9 @@ def visit_FunctionDef(self, node: FunctionDef) -> None:
                     if hasattr(argument, 'annotation') and argument.annotation:
                         self.visit(argument.annotation)
 
+            if node.returns:
+                self.visit(node.returns)
+
     def visit_AsyncFunctionDef(self, node: AsyncFunctionDef) -> None:
         """Remove and map function arguments and returns."""
         if self._function_is_wrapped_by_validate_arguments(node):
@@ -299,6 +302,9 @@ def visit_AsyncFunctionDef(self, node: AsyncFunctionDef) -> None:
                     if hasattr(argument, 'annotation') and argument.annotation:
                         self.visit(argument.annotation)
 
+            if node.returns:
+                self.visit(node.returns)
+
 
 class SQLAlchemyAnnotationVisitor(AnnotationVisitor):
     """Adds any names in the annotation to mapped names."""
@@ -548,24 +554,13 @@ def handle_fastapi_decorator(self, node: AsyncFunctionDef | FunctionDef) -> None
 
         To achieve this, we just visit the annotations to register them as "uses".
         """
-        for path in [node.args.args, node.args.kwonlyargs]:
+        for path in [node.args.args, node.args.kwonlyargs, node.args.posonlyargs]:
             for argument in path:
                 if hasattr(argument, 'annotation') and argument.annotation:
                     self.visit(argument.annotation)
-        if (
-            hasattr(node.args, 'kwarg')
-            and node.args.kwarg
-            and hasattr(node.args.kwarg, 'annotation')
-            and node.args.kwarg.annotation
-        ):
-            self.visit(node.args.kwarg.annotation)
-        if (
-            hasattr(node.args, 'vararg')
-            and node.args.vararg
-            and hasattr(node.args.vararg, 'annotation')
-            and node.args.vararg.annotation
-        ):
-            self.visit(node.args.vararg.annotation)
+
+        if node.returns:
+            self.visit(node.returns)
 
 
 class FunctoolsSingledispatchMixin:
@@ -627,6 +622,9 @@ def handle_singledispatch_decorator(self, node: FunctionDef | AsyncFunctionDef)
                 if hasattr(argument, 'annotation') and argument.annotation:
                     self.visit(argument.annotation)
 
+        if node.returns:
+            self.visit(node.returns)
+
 
 @dataclass
 class ImportName:

diff --git a/tests/test_fastapi_decorators.py b/tests/test_fastapi_decorators.py
@@ -9,7 +9,6 @@
 
 import pytest
 
-from flake8_type_checking.constants import TC002
 from tests.conftest import _get_error
 
 defaults = {'type_checking_fastapi_enabled': True}
@@ -56,9 +55,7 @@ def test_api_router_decorated_function_return_type(fdef):
             return None
         '''
     )
-    assert _get_error(example, error_code_filter='TC001,TC002,TC003', **defaults) == {
-        '5:0 ' + TC002.format(module='app.types.CustomType')
-    }
+    assert _get_error(example, error_code_filter='TC001,TC002,TC003', **defaults) == set()
 
 
 @pytest.mark.parametrize('fdef', ['def', 'async def'])