Skip to content

Commit

Permalink
feat: support multi paragraph sections inside the TYPE_CHECKING block
Browse files Browse the repository at this point in the history
  • Loading branch information
lyz-code committed Dec 21, 2021
1 parent 4572b4b commit 663bd4d
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 1 deletion.
3 changes: 2 additions & 1 deletion src/autoimport/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
}

Expand Down Expand Up @@ -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)

Expand Down
28 changes: 28 additions & 0 deletions tests/unit/test_services.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down

0 comments on commit 663bd4d

Please sign in to comment.