avoid double parse when removing imports
@@ -264,8 +264,8 @@ class FromImport(NamedTuple):
names: tuple[int, ...]
end: int
-
-def _parse_from_import(i: int, tokens: list[Token]) -> FromImport:
+ @classmethod
+ def parse(cls, i: int, tokens: list[Token]) -> FromImport:
j = i + 1
# XXX: does not handle explicit relative imports
while tokens[j].name != 'NAME':
@@ -290,22 +290,17 @@ def _parse_from_import(i: int, tokens: list[Token]) -> FromImport:
if tokens[names[i]].src == 'as':
del names[i:i + 2]
- return FromImport(mod_start, mod_end + 1, tuple(names), end)
-
+ return cls(mod_start, mod_end + 1, tuple(names), end)
-def _remove_import_partial(
- i: int,
- tokens: list[Token],
- *,
- idxs: list[int],
-) -> None:
- parsed = _parse_from_import(i, tokens)
+ def replace_modname(self, tokens: list[Token], modname: str) -> None:
+ tokens[self.mod_start:self.mod_end] = [Token('CODE', modname)]
+ def remove_parts(self, tokens: list[Token], idxs: list[int]) -> None:
for idx in reversed(idxs):
if idx == 0: # look forward until next name and del
- del tokens[parsed.names[idx]:parsed.names[idx + 1]]
+ del tokens[self.names[idx]:self.names[idx + 1]]
else: # look backward for comma and del
- j = end = parsed.names[idx]
+ j = end = self.names[idx]
while tokens[j].src != ',':
j -= 1
del tokens[j:end + 1]
@@ -324,8 +319,7 @@ def _replace_from_modname(
*,
modname: str,
) -> None:
- parsed = _parse_from_import(i, tokens)
- tokens[parsed.mod_start:parsed.mod_end] = [Token('CODE', modname)]
+ FromImport.parse(i, tokens).replace_modname(tokens, modname)
def _replace_from_mixed(
@@ -341,7 +335,7 @@ def _replace_from_mixed(
except ValueError:
return
- parsed = _parse_from_import(i, tokens)
+ parsed = FromImport.parse(i, tokens)
added_from_imports = collections.defaultdict(list)
for idx, mod, alias in exact_moves:
@@ -373,7 +367,7 @@ def _replace_from_mixed(
if len(parsed.names) == len(removal_idxs):
del tokens[i:parsed.end]
else:
- _remove_import_partial(i, tokens, idxs=removal_idxs)
+ parsed.remove_parts(tokens, removal_idxs)
@register(ast.ImportFrom)
