Repository: snok/flake8-type-checking (language: python)

Issue:
Improve type checking block discovery
Closes #10 


Discussion / hints:

Fix patch:
diff --git a/flake8_typing_only_imports/__init__.py b/flake8_typing_only_imports/__init__.py
index 1a6bc7e..9fc03e6 100644
--- a/flake8_typing_only_imports/__init__.py
+++ b/flake8_typing_only_imports/__init__.py
@@ -14,7 +14,7 @@ class Plugin:
     """Flake8 plugin."""
 
     name = 'flake8-typing-only-imports'
-    version = '0.1.6'
+    version = '0.1.7'
 
     def __init__(self, tree: 'Module') -> None:
         self.tree = tree
diff --git a/flake8_typing_only_imports/checker.py b/flake8_typing_only_imports/checker.py
index dfd79e9..33c4692 100644
--- a/flake8_typing_only_imports/checker.py
+++ b/flake8_typing_only_imports/checker.py
@@ -2,7 +2,7 @@
 import os
 from importlib.util import find_spec
 from pathlib import Path
-from typing import Any, Dict, Generator, List, Set, Tuple, Union
+from typing import Any, Dict, Generator, List, Optional, Set, Tuple, Union
 
 from flake8_typing_only_imports.constants import TYO100, TYO101, TYO200
 
@@ -31,6 +31,8 @@ def __init__(self, cwd: Path) -> None:
         self.type_checking_block_imports: Set[str] = set()
         self.unwrapped_annotations: List[Tuple[int, int, str]] = []
 
+        self.type_checking: Optional[Tuple[int, int]] = None
+
     # Map imports
 
     def _import_is_local(self, import_name: str) -> bool:
@@ -68,16 +70,31 @@ def _import_is_local(self, import_name: str) -> bool:
         origin = Path(spec.origin)
         return self.cwd in origin.parents
 
+    def _import_defined_inside_type_checking_block(self, node: Union[ast.Import, ast.ImportFrom]) -> bool:
+        """Indicate whether an import is defined inside an `if TYPE_CHECKING` block or not."""
+        if node.col_offset == 0:
+            return False
+        if self.type_checking is None:
+            return False
+        return self.type_checking[0] <= node.lineno <= self.type_checking[1]
+
+    def visit_If(self, node: ast.If) -> Any:
+        """Look for a TYPE_CHECKING block."""
+        if hasattr(node.test, 'id') and node.test.id == 'TYPE_CHECKING':  # type: ignore
+            self.type_checking = (node.lineno, node.end_lineno or node.lineno)
+        self.generic_visit(node)
+        return node
+
     def _add_import(self, node: Union[ast.Import, ast.ImportFrom]) -> None:
         """
         Add relevant ast objects to import lists.
 
         :param node: ast.Import or ast.ImportFrom object
-        :param names:  the string value of the code being imported
         """
-        if node.col_offset != 0:
-            # Avoid recording imports that live inside a `if TYPE_CHECKING` block
-            # The current handling is probably too naïve and could be upgraded
+        if self._import_defined_inside_type_checking_block(node):
+            # For type checking blocks we want to
+            # 1. Record annotations for TYO2XX errors
+            # 2. Avoid recording imports for TYO1XX errors, by returning early
             for name_node in node.names:
                 if hasattr(name_node, 'asname') and name_node.asname:
                     name = name_node.asname
@@ -128,11 +145,8 @@ def visit_Name(self, node: ast.Name) -> ast.Name:
 
     def visit_AnnAssign(self, node: ast.AnnAssign) -> ast.AnnAssign:
         """Remove all annotation assignments."""
-        if hasattr(node, 'annotation') and node.annotation:
-            if isinstance(node.annotation, ast.Name):
-                self.unwrapped_annotations.append((node.lineno, node.col_offset, node.annotation.id))
-            else:
-                print('UNHANDLED TYPE:', type(node.annotation))  # noqa  # todo: remove in a future iteration
+        if hasattr(node, 'annotation') and node.annotation and isinstance(node.annotation, ast.Name):
+            self.unwrapped_annotations.append((node.lineno, node.col_offset, node.annotation.id))
             delattr(node, 'annotation')
         self.generic_visit(node)
         return node
diff --git a/pyproject.toml b/pyproject.toml
index 1ebc515..5ec9ca0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "flake8-typing-only-imports"
-version = "0.1.6"  # remember to also change __init__.py
+version = "0.1.7"  # remember to also change __init__.py
 description = "A flake8 plugin that flags imports exclusively used for type annotations"
 homepage = "https://github.com/sondrelg/flake8-typing-only-imports"
 repository = "https://github.com/sondrelg/flake8-typing-only-imports"
@@ -10,6 +10,7 @@ license = "BSD-2-Clause"
 readme = "README.md"
 keywords = ["flake8", "lint", "type hint", "typing", "type annotations", "imports"]
 classifiers = [
+    "Development Status :: 2 - Pre-Alpha",
     "Intended Audience :: Developers",
     "Environment :: Console",
     "Operating System :: OS Independent",


Test patch:
diff --git a/tests/test_tyo200.py b/tests/test_tyo200.py
index ddf3794..56d58b5 100644
--- a/tests/test_tyo200.py
+++ b/tests/test_tyo200.py
@@ -29,6 +29,27 @@
             '4:12 ' + TYO200.format(annotation='something'),
         },
     ),
+    (
+        textwrap.dedent(
+            """
+        def test():
+            import x
+            y = x
+        """
+        ),
+        set(),
+    ),
+    (
+        textwrap.dedent(
+            """
+        if TYPE_CHECKING:
+            import x
+
+        y: x
+        """
+        ),
+        {'5:0 ' + TYO200.format(annotation='x')},
+    ),
 ]