diff --git a/pyupgrade/_plugins/imports.py b/pyupgrade/_plugins/imports.py index 2883f958..8a11ca41 100644 --- a/pyupgrade/_plugins/imports.py +++ b/pyupgrade/_plugins/imports.py @@ -10,7 +10,6 @@ from tokenize_rt import Offset from tokenize_rt import Token -from tokenize_rt import UNIMPORTANT_WS from pyupgrade._ast_helpers import ast_to_offset from pyupgrade._data import register @@ -18,6 +17,7 @@ from pyupgrade._data import TokenFunc from pyupgrade._token_helpers import find_end from pyupgrade._token_helpers import find_token +from pyupgrade._token_helpers import indented_amount # GENERATED VIA generate-imports # Using reorder-python-imports==3.6.0 @@ -335,17 +335,10 @@ def _replace_from_mixed( exact_moves: list[tuple[int, str, ast.alias]], module_moves: list[tuple[int, str, ast.alias]], ) -> None: - if i == 0: - indent = '' - elif tokens[i - 1].name in {UNIMPORTANT_WS, 'INDENT'}: - if tokens[i - 2].name in {'NL', 'NEWLINE', 'DEDENT'}: - indent = tokens[i - 1].src - else: # inline import - return - elif tokens[i - 1].name not in {'NL', 'NEWLINE', 'DEDENT'}: - return # inline import - else: - indent = '' + try: + indent = indented_amount(i, tokens) + except ValueError: + return parsed = _parse_from_import(i, tokens) @@ -361,7 +354,7 @@ def _replace_from_mixed( if new_mod: added_from_imports[new_mod].append(_alias_to_s(new_alias)) else: - new_imports.append(f'import {_alias_to_s(new_alias)}\n') + new_imports.append(f'{indent}import {_alias_to_s(new_alias)}\n') bisect.insort(removal_idxs, idx) new_imports.extend( @@ -450,6 +443,11 @@ def _replace_import( exact_moves: list[tuple[int, str, ast.alias]], to_from: list[tuple[int, str, str, ast.alias]], ) -> None: + try: + indent = indented_amount(i, tokens) + except ValueError: + return + end = find_end(tokens, i) parts = [] @@ -474,7 +472,7 @@ def _replace_import( tokens[slice(*parts[idx])] = [Token('CODE', _alias_to_s(new_alias))] new_imports = sorted( - f'from {new_mod} import ' + f'{indent}from {new_mod} import ' f'{_alias_to_s(ast.alias(name=new_sym, asname=alias.asname))}\n' for _, new_mod, new_sym, alias in to_from ) diff --git a/pyupgrade/_token_helpers.py b/pyupgrade/_token_helpers.py index 36e7f796..8a8e916f 100644 --- a/pyupgrade/_token_helpers.py +++ b/pyupgrade/_token_helpers.py @@ -515,3 +515,17 @@ def replace_list_comp_brackets(i: int, tokens: list[Token]) -> None: end = find_closing_bracket(tokens, start) tokens[end] = Token('OP', ')') tokens[start] = Token('OP', '(') + + +def indented_amount(i: int, tokens: list[Token]) -> str: + if i == 0: + return '' + elif i >= 2 and tokens[i - 1].name in {UNIMPORTANT_WS, 'INDENT'}: + if tokens[i - 2].name in {'NL', 'NEWLINE', 'DEDENT'}: + return tokens[i - 1].src + else: # inline import + raise ValueError('not at beginning of line') + elif tokens[i - 1].name not in {'NL', 'NEWLINE', 'DEDENT'}: + raise ValueError('not at beginning of line') + else: + return '' diff --git a/tests/features/import_replaces_test.py b/tests/features/import_replaces_test.py index 5f3392c8..d53acbe0 100644 --- a/tests/features/import_replaces_test.py +++ b/tests/features/import_replaces_test.py @@ -23,12 +23,17 @@ pytest.param( 'if True: from six.moves import getcwd, StringIO\n', (3,), - id='play stupid games, win stupid prizes pt1', + id='inline from-import with space', ), pytest.param( 'if True:from six.moves import getcwd, StringIO\n', (3,), - id='play stupid games, win stupid prizes pt2', + id='inline from-import without space', + ), + pytest.param( + 'if True:import mock, sys\n', + (3,), + id='inline import-import', ), pytest.param( 'import xml.etree.cElementTree', @@ -136,7 +141,16 @@ def test_mock_noop_keep_mock(): 'if True:\n' ' from collections import Counter\n' ' from collections.abc import Mapping\n', - id='indented import being added', + id='indented from-import being added', + ), + pytest.param( + 'if True:\n' + ' from six.moves import queue, urllib_request\n', + (3,), + 'if True:\n' + ' from six.moves import urllib_request\n' + ' import queue\n', + id='indented import-import being added', ), pytest.param( 'if True:\n'