Merge pull request #173 from lyz-code/feat/support_multiparagraph_typ…

…ing_imports

feat: support multi paragraph sections inside the TYPE_CHECKING block

src/autoimport/model.py

@@ -33,6 +33,7 @@
     "StringIO": "from io import StringIO",
     "suppress": "from contextlib import suppress",
     "TempdirFactory": "from _pytest.tmpdir import TempdirFactory",
+    "tz": "from dateutil import tz",
     "YAMLError": "from yaml import YAMLError",
 }
 
@@ -165,7 +166,7 @@ def _extract_typing_statements(self, source_lines: List[str]) -> None:
             self.typing.append(source_lines[typing_start_line])
             typing_start_line += 1
             for line in source_lines[typing_start_line:]:
-                if not re.match(r"^\s+.*", line):
+                if not re.match(r"^\s+.*", line) and line != "":
                     break
                 self.typing.append(line)
 

tests/unit/test_services.py

@@ -712,6 +712,34 @@ def read_book(book: Book):
     assert result == source
 
 
+def test_fix_respects_multiparagraph_type_checking_import_statements() -> None:
+    """
+    Given: Code with two paragraphs of imports inside an if TYPE_CHECKING block
+    When: Fix code is run.
+    Then: The imports are not moved above the if statement.
+    """
+    source = dedent(
+        """\
+        import os
+        from typing import TYPE_CHECKING
+
+        if TYPE_CHECKING:
+            from .model import Book
+
+            from other import Other
+
+        os.getcwd()
+
+
+        def read_book(book: Book, other: Other):
+            pass"""
+    )
+
+    result = fix_code(source)
+
+    assert result == source
+
+
 def test_fix_respects_try_except_in_import_statements() -> None:
     """
     Given: Code with try except statements in the imports.