Repository: snok/flake8-type-checking

Issue:
False positive for FastAPI return types.
As mentioned in #52 

According to this test, `CustomType` should be moved to a `TYPE_CHECKING` block.

```python
@pytest.mark.parametrize('fdef', ['def', 'async def'])
def test_api_router_decorated_function_return_type(fdef):
    """
    We don't care about return types. To my knowledge,
    these are not evaluated by FastAPI/pydantic.
    """
    example = textwrap.dedent(f'''
        from fastapi import APIRouter
        from fastapi import Request

        from app.types import CustomType

        some_router = APIRouter(prefix='/some-path')

        @some_router.get('/{{resource_id}}')
        {fdef} list_something(request: Request) -> CustomType:
            return None
        ''')
    assert _get_error(example, error_code_filter='TC001,TC002,TC003', **defaults) == {
        '5:0 ' + TC002.format(module='app.types.CustomType')
    }
```

However, doing so, even with a pretty basic route gives the following error in FastAPI:

```python
     def test() -> Test:
                   ^^^^
 NameError: name 'Test' is not defined
```

Hints from issue comments:
@sondrelg: Agree that this seems unexpected. Looking at this quickly, I would expect:

```python
from fastapi import APIRouter
from fastapi import Request

from app.types import CustomType

some_router = APIRouter(prefix='/some-path')


@some_router.get('/{{resource_id}}')
def list_something(request: Request) -> "CustomType":
    return None
```

or 

```python
from __future__ import annotations

from fastapi import APIRouter
from fastapi import Request

from app.types import CustomType

some_router = APIRouter(prefix='/some-path')


@some_router.get('/{{resource_id}}')
def list_something(request: Request) -> CustomType:
    return None
```

To be correct, wrt. to the return type.

Would you be interested in creating a fix for this @danpker? I'm on holiday at the moment, so won't have time to look at this, but will make time to push a fix if one is created 👍 

Solution patch:
diff --git a/flake8_type_checking/checker.py b/flake8_type_checking/checker.py
index 93d28df..d866926 100644
--- a/flake8_type_checking/checker.py
+++ b/flake8_type_checking/checker.py
@@ -291,6 +291,9 @@ class PydanticMixin:
                     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 @@ class PydanticMixin:
                     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 @@ class FastAPIMixin:
 
         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 @@ class FunctoolsSingledispatchMixin:
                 if hasattr(argument, 'annotation') and argument.annotation:
                     self.visit(argument.annotation)
 
+        if node.returns:
+            self.visit(node.returns)
+
 
 @dataclass
 class ImportName:


Test patch:
diff --git a/tests/test_fastapi_decorators.py b/tests/test_fastapi_decorators.py
index a47746d..5ae158a 100644
--- a/tests/test_fastapi_decorators.py
+++ b/tests/test_fastapi_decorators.py
@@ -9,7 +9,6 @@ import textwrap
 
 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'])