Add linting for imports under if TYPE_CHECKING: block (#205)

* Add linting for imports under `if TYPE_CHECKING:` block

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Ian Stapleton Cordasco <graffatcolmingov@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

README.rst

@@ -73,7 +73,8 @@ Limitations
 -----------
 
 Currently these checks are limited to module scope imports only.
-Conditional imports in module scope will also be ignored.
+Conditional imports in module scope will be ignored except imports
+under ```if TYPE_CHECKING:``` block.
 
 Classification of an imported module is achieved by checking the
 module against a stdlib list and then if there is no match against the

flake8_import_order/__init__.py

@@ -32,7 +32,16 @@
 
 ClassifiedImport = namedtuple(
     "ClassifiedImport",
-    ["type", "is_from", "modules", "names", "lineno", "level", "package"],
+    [
+        "type",
+        "is_from",
+        "modules",
+        "names",
+        "lineno",
+        "level",
+        "package",
+        "type_checking",
+    ],
 )
 NewLine = namedtuple("NewLine", ["lineno"])
 
@@ -85,8 +94,13 @@ def __init__(self, application_import_names, application_package_names):
         self.application_import_names = frozenset(application_import_names)
         self.application_package_names = frozenset(application_package_names)
 
+    def generic_visit(self, node):
+        for child in ast.iter_child_nodes(node):
+            child.parent = node
+        return super().generic_visit(node)
+
     def visit_Import(self, node):  # noqa: N802
-        if node.col_offset == 0:
+        if node.col_offset == 0 or self._type_checking_import(node):
             modules = [alias.name for alias in node.names]
             types_ = {self._classify_type(module) for module in modules}
             if len(types_) == 1:
@@ -101,11 +115,12 @@ def visit_Import(self, node):  # noqa: N802
                 node.lineno,
                 0,
                 root_package_name(modules[0]),
+                self._type_checking_import(node),
             )
             self.imports.append(classified_import)
 
     def visit_ImportFrom(self, node):  # noqa: N802
-        if node.col_offset == 0:
+        if node.col_offset == 0 or self._type_checking_import(node):
             module = node.module or ""
             if node.level > 0:
                 type_ = ImportType.APPLICATION_RELATIVE
@@ -120,9 +135,16 @@ def visit_ImportFrom(self, node):  # noqa: N802
                 node.lineno,
                 node.level,
                 root_package_name(module),
+                self._type_checking_import(node),
             )
             self.imports.append(classified_import)
 
+    def _type_checking_import(self, node):
+        return (
+            isinstance(node.parent, ast.If)
+            and node.parent.test.id == "TYPE_CHECKING"
+        )
+
     def _classify_type(self, module):
         package_names = get_package_names(module)
 

flake8_import_order/styles.py

@@ -42,6 +42,13 @@ def check(self):
     def _check(self, previous_import, previous, current_import):
         yield from self._check_I666(current_import)
         yield from self._check_I101(current_import)
+        if (
+            previous_import is not None
+            and not previous_import.type_checking
+            and current_import.type_checking
+        ):
+            yield from self._check_I300(previous_import, current_import)
+            previous_import = None
         if previous_import is not None:
             yield from self._check_I100(previous_import, current_import)
             yield from self._check_I201(
@@ -70,6 +77,14 @@ def _check_I101(self, current_import):  # noqa: N802
                 "Should be {}".format(corrected),
             )
 
+    def _check_I300(self, previous_import, current_import):  # noqa: N802
+        if current_import.lineno - previous_import.lineno != 3:
+            yield Error(
+                current_import.lineno,
+                "I300",
+                "TYPE_CHECKING block should have one newline above.",
+            )
+
     def _check_I100(self, previous_import, current_import):  # noqa: N802
         previous_key = self.import_key(previous_import)
         current_key = self.import_key(current_import)

tests/test_cases/additional_newline.py

@@ -22,3 +22,28 @@
 from . import A
 
 from . import B  # I202
+
+if TYPE_CHECKING:
+    import ast
+    # This comment should not trigger a I202 (not a newline)
+    import os
+
+    import signal # I202
+
+    import X
+    from X import B, b, \
+        C, d
+
+    from Y import A # I202
+    from Y import (
+        B, b,
+        C, d,
+    )
+    from Z import A
+
+    import flake8_import_order
+
+    import tests  # I202
+    from . import A
+
+    from . import B  # I202

tests/test_cases/additional_newline_cryptography.py

@@ -14,3 +14,20 @@
 from . import A
 
 from . import B  # I202
+
+if TYPE_CHECKING:
+    import ast
+
+    import signal # I202
+
+    import X
+
+    import Y
+
+    import flake8_import_order
+
+    import tests
+
+    from . import A
+
+    from . import B  # I202

tests/test_cases/additional_newline_edited.py

@@ -23,3 +23,29 @@
 from . import A
 
 from . import B  # I202
+
+if TYPE_CHECKING:
+    import ast
+    # This comment should not trigger a I202 (not a newline)
+    import os
+
+    import signal # I202
+
+    import X
+    from X import B, b, \
+        C, d
+
+    from Y import A # I202
+    from Y import (
+        B, b,
+        C, d,
+    )
+    from Z import A
+
+    import flake8_import_order
+
+    import tests  # I202
+
+    from . import A
+
+    from . import B  # I202

tests/test_cases/complete_appnexus.py

@@ -32,3 +32,38 @@
 from .. import B
 from ..A import A
 from ..B import B
+
+if TYPE_CHECKING:
+    import ast
+    from functools import *
+    import os
+    from os import path
+    import StringIO
+    import sys
+
+    import X
+    from X import *
+    from X import A
+    from X import B, b, C, d
+    import Y
+    from Y import *
+    from Y import A
+    from Y import B, C, D
+    from Y import e
+    import Z
+    from Z import A
+    from Z.A import A
+    from Z.A.B import A
+
+    from localpackage import A, b
+
+    import flake8_import_order
+    from flake8_import_order import *
+    from . import A
+    from . import B
+    from .A import A
+    from .B import B
+    from .. import A
+    from .. import B
+    from ..A import A
+    from ..B import B

tests/test_cases/complete_cryptography.py

@@ -40,3 +40,45 @@
 from .. import B
 from ..A import A
 from ..B import B
+
+if TYPE_CHECKING:
+    import ast
+    import os
+    import sys
+    from functools import *
+    from os import path
+
+    import X
+    from X import *
+    from X import A
+    from X import B, C, D
+
+    import Y
+    from Y import *
+    from Y import A
+    from Y import B, C, D
+
+    import Z
+    from Z import A
+    from Z.A import A
+    from Z.A.B import A
+
+    import localpackage
+
+    import flake8_import_order
+    from flake8_import_order import *
+    from flake8_import_order import A
+    from flake8_import_order import B
+
+    import tests
+    from tests import A
+    from tests import B
+
+    from . import A
+    from . import B
+    from .A import A
+    from .B import B
+    from .. import A
+    from .. import B
+    from ..A import A
+    from ..B import B

tests/test_cases/complete_edited.py

@@ -35,3 +35,41 @@
 from .. import B
 from ..A import A
 from ..B import B
+
+if TYPE_CHECKING:
+    import ast
+    import os
+    import StringIO
+    import sys
+    from functools import *
+    from os import path
+
+    import X
+    import Y
+    import Z
+    from X import *
+    from X import A
+    from X import B, b, C, d
+    from Y import *
+    from Y import A
+    from Y import B, C, D
+    from Y import e
+    from Z import A
+    from Z.A import A
+    from Z.A.B import A
+
+    import localpackage
+
+    from localpackage import A, b
+
+    import flake8_import_order
+    from flake8_import_order import *
+
+    from . import A
+    from . import B
+    from .A import A
+    from .B import B
+    from .. import A
+    from .. import B
+    from ..A import A
+    from ..B import B

tests/test_cases/complete_google.py

@@ -31,3 +31,37 @@
 from .. import B
 from ..A import A
 from ..B import B
+
+if TYPE_CHECKING:
+    import ast
+    from functools import *
+    import os
+    from os import path
+    import StringIO
+    import sys
+
+    import localpackage
+    import X
+    from X import *
+    from X import A
+    from X import B, b, C, d
+    import Y
+    from Y import *
+    from Y import A
+    from Y import B, C, D
+    from Y import e
+    import Z
+    from Z import A
+    from Z.A import A
+    from Z.A.B import A
+
+    import flake8_import_order
+    from flake8_import_order import *
+    from . import A
+    from . import B
+    from .A import A
+    from .B import B
+    from .. import A
+    from .. import B
+    from ..A import A
+    from ..B import B