Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: support multi paragraph sections inside the TYPE_CHECKING block #173

Merged
merged 1 commit into from Dec 21, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
3 changes: 2 additions & 1 deletion src/autoimport/model.py
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
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