Add linting for imports under if TYPE_CHECKING: block

Author: birthdaysgift

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

@@ -17,7 +17,8 @@
 
 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'])
 
@@ -70,8 +71,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:
@@ -81,11 +87,12 @@ def visit_Import(self, node):  # noqa: N802
             classified_import = ClassifiedImport(
                 type_, False, modules, [], 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
@@ -96,9 +103,16 @@ def visit_ImportFrom(self, node):  # noqa: N802
                 type_, True, [module], names,
                 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

@@ -37,6 +37,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(previous_import, previous, current_import)
@@ -61,6 +68,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

@@ -33,3 +33,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

@@ -36,3 +36,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

@@ -32,3 +32,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

tests/test_cases/complete_pep8.py

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