feat: Add support for singledispatch/singledispatchmethod

Author: Daverball

flake8_type_checking/checker.py

@@ -582,6 +582,66 @@ def handle_fastapi_decorator(self, node: Union[AsyncFunctionDef, FunctionDef]) -
             self.visit(node.args.vararg.annotation)
 
 
+class FunctoolsSingledispatchMixin:
+    """
+    Contains the necessary logic for `functools.singledispatch` support.
+
+    `functools.singledispatch` and `functools.singledispatchmethod` require
+    runtime access to all annotations.
+
+    ```python
+    from functools import singledispatch
+
+    from mylib import Foo
+
+    @singledispatch
+    def foo(arg: Foo) -> str:
+        return arg.name
+    ```
+
+    Since the only use of `Foo` is within an annotation, we would usually emit
+    a TC003 for the `mylib` import. But since `singledispatch` requires runtime
+    access to `Foo`, that would be a false positive.
+    """
+
+    if TYPE_CHECKING:
+
+        def in_type_checking_block(self, lineno: int, col_offset: int) -> bool:  # noqa: D102
+            ...
+
+        def lookup_full_name(self, node: ast.AST) -> str | None:  # noqa: D102
+            ...
+
+        def visit(self, node: ast.AST) -> ast.AST:  # noqa: D102
+            ...
+
+    def visit_FunctionDef(self, node: FunctionDef) -> None:
+        """Remove and map function arguments and returns."""
+        super().visit_FunctionDef(node)  # type: ignore[misc]
+        if self.has_singledispatch_decorator(node):
+            self.handle_singledispatch_decorator(node)
+
+    def visit_AsyncFunctionDef(self, node: AsyncFunctionDef) -> None:
+        """Remove and map function arguments and returns."""
+        super().visit_AsyncFunctionDef(node)  # type: ignore[misc]
+        if self.has_singledispatch_decorator(node):
+            self.handle_singledispatch_decorator(node)
+
+    def has_singledispatch_decorator(self, node: FunctionDef | AsyncFunctionDef) -> bool:
+        """Determine whether this function is decorated with `functools.singledispatch`."""
+        return any(
+            self.lookup_full_name(decorator_node) in ('functools.singledispatch', 'functools.singledispatchmethod')
+            for decorator_node in node.decorator_list
+        )
+
+    def handle_singledispatch_decorator(self, node: FunctionDef | AsyncFunctionDef) -> None:
+        """Walk all the annotations to register them as runtime uses."""
+        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)
+
+
 @dataclass
 class ImportName:
     """DTO for representing an import in different string-formats."""
@@ -949,7 +1009,14 @@ def visit_annotated_value(self, node: ast.expr) -> None:
 
 
 class ImportVisitor(
-    DunderAllMixin, AttrsMixin, InjectorMixin, FastAPIMixin, PydanticMixin, SQLAlchemyMixin, ast.NodeVisitor
+    DunderAllMixin,
+    FunctoolsSingledispatchMixin,
+    AttrsMixin,
+    InjectorMixin,
+    FastAPIMixin,
+    PydanticMixin,
+    SQLAlchemyMixin,
+    ast.NodeVisitor,
 ):
     """Map all imports outside of type-checking blocks."""
 

tests/test_tc001_to_tc003.py

@@ -120,6 +120,47 @@ def get_tc_001_to_003_tests(import_: str, ERROR: str) -> L:
             (f'import {import_}\ntype x = {import_}', {f"1:0 {ERROR.format(module=f'{import_}')}"})
         )
 
+    # Imports used for `functools.singledispatch`. None of these should generate errors.
+    used_for_singledispatch: L = [
+        (
+            textwrap.dedent(f'''
+                import functools
+
+                from {import_} import Dict, Any
+
+                @functools.singledispatch
+                def foo(arg: Dict[str, Any]) -> Any:
+                    return 1
+                '''),
+            set(),
+        ),
+        (
+            textwrap.dedent(f'''
+                from functools import singledispatch
+
+                from {import_} import Dict, Any
+
+                @singledispatch
+                def foo(arg: Dict[str, Any]) -> Any:
+                    return 1
+                '''),
+            set(),
+        ),
+        (
+            textwrap.dedent(f'''
+                from functools import singledispatchmethod
+
+                from {import_} import Dict, Any
+
+                class Foo:
+                    @singledispatchmethod
+                    def foo(self, arg: Dict[str, Any]) -> Any:
+                        return 1
+                '''),
+            set(),
+        ),
+    ]
+
     other_useful_test_cases: L = [
         (
             textwrap.dedent(f'''
@@ -237,6 +278,7 @@ class Migration:
         *used_for_arg_annotations_only,
         *used_for_return_annotations_only,
         *used_for_type_alias_only,
+        *used_for_singledispatch,
         *other_useful_test_cases,
     ]
 

tests/test_tc004.py

@@ -227,6 +227,47 @@ class X:
             '3:0 ' + TC004.format(module='z'),
         },
     ),
+    # functools.singledispatch
+    (
+        textwrap.dedent("""
+            import functools
+
+            if TYPE_CHECKING:
+                from foo import FooType
+
+            @functools.singledispatch
+            def foo(arg: FooType) -> int:
+                return 1
+            """),
+        {'5:0 ' + TC004.format(module='FooType')},
+    ),
+    (
+        textwrap.dedent("""
+            from functools import singledispatch
+
+            if TYPE_CHECKING:
+                from foo import FooType
+
+            @functools.singledispatch
+            def foo(arg: FooType) -> int:
+                return 1
+            """),
+        {'5:0 ' + TC004.format(module='FooType')},
+    ),
+    (
+        textwrap.dedent("""
+            from functools import singledispatchmethod
+
+            if TYPE_CHECKING:
+                from foo import FooType
+
+            class Foo:
+                @functools.singledispatch
+                def foo(self, arg: FooType) -> int:
+                    return 1
+            """),
+        {'5:0 ' + TC004.format(module='FooType')},
+    ),
 ]