gh-69605: PyREPL: insert "import" after "from foo <tab>" (#148445)

Improve _pyrepl._module_completer.ModuleCompleter.get_completions:
  * "from x.y.z " [tab]-> "from x.y.z import "
  * "from foo"    [tab]-> "from foo "          [tab]-> "from foo import"
    (if only one suggestion)

Co-authored-by: Lukasz Langa <ambv@meta.com>
This commit is contained in:
Loïc SimonandLukasz Langa authored and GitHub committed 2026-10-21 03:14:55 -07:00
1 parent 7ed37e95c5
commit fdee11f099
3 files changed
+88 -65

No files matched your search

+34 -20
View File
@@ -97,6 +97,7 @@ class ModuleCompleter:
self,
from_name: str | None,
name: str | None,
space_end: bool,
*,
include_values: bool = True,
) -> tuple[list[str], list[Any], CompletionAction | None]:
@@ -111,10 +112,16 @@ class ModuleCompleter:
return names, values, None
if name is None:
if space_end and from_name:
# from x.y.z <tab>
return ["import "], [sys] if include_values else [], None
# from x.y.z<tab>
path, prefix = self.get_path_and_prefix(from_name)
modules = self.find_modules(path, prefix)
names = [self.format_completion(path, module) for module in modules]
if names == [from_name]:
# Exact match already written: continue with import statement
names = [f"{from_name} import "]
# These are always modules, use dummy values to get the right color
values = [sys] * len(names) if include_values else []
return names, values, None
@@ -367,11 +374,12 @@ class ImportParser:
suitable for autocomplete suggestions.
Examples:
- import foo -> Result(from_name=None, name='foo')
- import foo. -> Result(from_name=None, name='foo.')
- from foo -> Result(from_name='foo', name=None)
- from foo import bar -> Result(from_name='foo', name='bar')
- from .foo import ( -> Result(from_name='.foo', name='')
- `import foo` -> Result(from_name=None, name='foo')
- `import foo.` -> Result(from_name=None, name='foo.')
- `from foo` -> Result(from_name='foo', name=None)
- `from foo ` -> Result(from_name='foo', name=None, space_end=True)
- `from foo import bar` -> Result(from_name='foo', name='bar')
- `from .foo import (` -> Result(from_name='.foo', name='')
Note that the parser works in reverse order, starting from the
last token in the input string. This makes the parser more robust
@@ -401,10 +409,10 @@ class ImportParser:
tokens = []
self.tokens = TokenQueue(tokens[::-1])
def parse(self) -> tuple[str | None, str | None] | None:
def parse(self) -> tuple[str | None, str | None, bool] | None:
if not (res := self._parse()):
return None
return res.from_name, res.name
return res.from_name, res.name, res.space_end
def _parse(self) -> Result | None:
with self.tokens.save_state():
@@ -414,7 +422,7 @@ class ImportParser:
def parse_import(self) -> Result:
if self.code.rstrip().endswith('import') and self.code.endswith(' '):
return Result(name='')
return Result(name='', space_end=True)
if self.tokens.peek_string(','):
name = ''
else:
@@ -427,27 +435,32 @@ class ImportParser:
self.tokens.pop()
self.parse_dotted_as_name()
if self.tokens.peek_string('import'):
return Result(name=name)
return Result(name=name, space_end=self.code.endswith(' '))
raise ParseError('parse_import')
def parse_from_import(self) -> Result:
space_end = self.code.endswith(' ')
stripped = self.code.rstrip()
if stripped.endswith('import') and self.code.endswith(' '):
return Result(from_name=self.parse_empty_from_import(), name='')
if stripped.endswith('from') and self.code.endswith(' '):
return Result(from_name='')
if stripped.endswith('import') and space_end:
from_name = self.parse_empty_from_import()
return Result(from_name=from_name, name='', space_end=space_end)
if stripped.endswith('from') and space_end:
return Result(from_name='', space_end=space_end)
if self.tokens.peek_string('(') or self.tokens.peek_string(','):
return Result(from_name=self.parse_empty_from_import(), name='')
if self.code.endswith(' '):
raise ParseError('parse_from_import')
from_name = self.parse_empty_from_import()
return Result(from_name=from_name, name='', space_end=space_end)
name = self.parse_dotted_name()
if '.' in name:
if name.endswith(".") and space_end:
raise ParseError('parse_from_import')
self.tokens.pop_string('from')
return Result(from_name=name)
return Result(from_name=name, space_end=space_end)
if self.tokens.peek_string('from'):
return Result(from_name=name)
return Result(from_name=name, space_end=space_end)
if space_end:
raise ParseError('parse_from_import')
from_name = self.parse_empty_from_import()
return Result(from_name=from_name, name=name)
return Result(from_name=from_name, name=name, space_end=space_end)
def parse_empty_from_import(self) -> str:
if self.tokens.peek_string(','):
@@ -513,10 +526,11 @@ class ParseError(Exception):
pass
@dataclass(frozen=True)
@dataclass(frozen=True, kw_only=True)
class Result:
from_name: str | None = None
name: str | None = None
space_end: bool
class TokenQueue:
+52 -45
View File
@@ -1231,12 +1231,20 @@ class TestPyReplModuleCompleter(TestCase):
("import foo, impo\t\n", "import foo, importlib"),
("import foo as bar, impo\t\n", "import foo as bar, importlib"),
("from impo\t\n", "from importlib"),
("from impo\t\t\n", "from importlib import "),
("from impo\t\t\t\n", "from importlib import "),
("from impo \t\n", "from impo import "),
("from importlib\t\n", "from importlib import "),
("from importlib.res\t\n", "from importlib.resources"),
("from importlib.\t\tres\t\n", "from importlib.resources"),
("from importlib.res\t\t\n", "from importlib.resources import "),
("from importlib.res \t\n", "from importlib.res import "),
("from importlib.resources\t\n", "from importlib.resources import "),
("from importlib.resources.ab\t\n", "from importlib.resources.abc"),
("from importlib import mac\t\n", "from importlib import machinery"),
("from importlib import res\t\n", "from importlib import resources"),
("from importlib.res\t import a\t\n", "from importlib.resources import abc"),
("from importlib.res\t\ta\t\n", "from importlib.resources import abc"),
("from __phello__ import s\t\n", "from __phello__ import spam"), # frozen module
("import importlib .res\t\n", "import importlib .resources"),
("import importlib. res\t\n", "import importlib. resources"),
@@ -1616,46 +1624,48 @@ class TestPyReplModuleCompleter(TestCase):
def test_parse(self):
cases = (
('import ', (None, '')),
('import foo', (None, 'foo')),
('import foo,', (None, '')),
('import foo, ', (None, '')),
('import foo, bar', (None, 'bar')),
('import foo, bar, baz', (None, 'baz')),
('import foo as bar,', (None, '')),
('import foo as bar, ', (None, '')),
('import foo as bar, baz', (None, 'baz')),
('import a.', (None, 'a.')),
('import a.b', (None, 'a.b')),
('import a.b.', (None, 'a.b.')),
('import a.b.c', (None, 'a.b.c')),
('import a.b.c, foo', (None, 'foo')),
('import a.b.c, foo.bar', (None, 'foo.bar')),
('import a.b.c, foo.bar,', (None, '')),
('import a.b.c, foo.bar, ', (None, '')),
('from foo', ('foo', None)),
('from a.', ('a.', None)),
('from a .', ('a.', None)),
('from a .', ('a.', None)),
('from a .', ('a.', None)),
('from a.b', ('a.b', None)),
('from a.b.', ('a.b.', None)),
('from a. b.', ('a.b.', None)),
('from a . b .', ('a.b.', None)),
('from a.b.c', ('a.b.c', None)),
('from foo import ', ('foo', '')),
('from foo import a', ('foo', 'a')),
('from ', ('', None)),
('from . import a', ('.', 'a')),
('from .foo import a', ('.foo', 'a')),
('from ..foo import a', ('..foo', 'a')),
('from foo import (', ('foo', '')),
('from foo import ( ', ('foo', '')),
('from foo import (a', ('foo', 'a')),
('from foo import (a,', ('foo', '')),
('from foo import (a, ', ('foo', '')),
('from foo import (a, c', ('foo', 'c')),
('from foo import (a as b, c', ('foo', 'c')),
('import ', (None, '', True)),
('import foo', (None, 'foo', False)),
('import foo,', (None, '', False)),
('import foo, ', (None, '', True)),
('import foo, bar', (None, 'bar', False)),
('import foo, bar, baz', (None, 'baz', False)),
('import foo as bar,', (None, '', False)),
('import foo as bar, ', (None, '', True)),
('import foo as bar, baz', (None, 'baz', False)),
('import a.', (None, 'a.', False)),
('import a.b', (None, 'a.b', False)),
('import a.b.', (None, 'a.b.', False)),
('import a.b.c', (None, 'a.b.c', False)),
('import a.b.c, foo', (None, 'foo', False)),
('import a.b.c, foo.bar', (None, 'foo.bar', False)),
('import a.b.c, foo.bar,', (None, '', False)),
('import a.b.c, foo.bar, ', (None, '', True)),
('from foo', ('foo', None, False)),
('from foo ', ('foo', None, True)),
('from a.', ('a.', None, False)),
('from a .', ('a.', None, False)),
('from a .', ('a.', None, False)),
('from a .', ('a.', None, False)),
('from a.b', ('a.b', None, False)),
('from a.b.', ('a.b.', None, False)),
('from a.b ', ('a.b', None, True)),
('from a. b.', ('a.b.', None, False)),
('from a . b .', ('a.b.', None, False)),
('from a.b.c', ('a.b.c', None, False)),
('from foo import ', ('foo', '', True)),
('from foo import a', ('foo', 'a', False)),
('from ', ('', None, True)),
('from . import a', ('.', 'a', False)),
('from .foo import a', ('.foo', 'a', False)),
('from ..foo import a', ('..foo', 'a', False)),
('from foo import (', ('foo', '', False)),
('from foo import ( ', ('foo', '', True)),
('from foo import (a', ('foo', 'a', False)),
('from foo import (a,', ('foo', '', False)),
('from foo import (a, ', ('foo', '', True)),
('from foo import (a, c', ('foo', 'c', False)),
('from foo import (a as b, c', ('foo', 'c', False)),
)
for code, parsed in cases:
parser = ImportParser(code)
@@ -1679,12 +1689,9 @@ class TestPyReplModuleCompleter(TestCase):
cases = (
'',
'import foo ',
'from foo ',
'import foo. ',
'import foo.bar ',
'from foo ',
'from foo. ',
'from foo.bar ',
'from foo import bar ',
'from foo import (bar ',
'from foo import bar, baz ',
@@ -1723,9 +1730,9 @@ class TestPyReplModuleCompleter(TestCase):
'if 1:\n pass\n\tpass', # _tokenize TabError -> tokenize TabError
)
for code in cases:
parser = ImportParser(code)
actual = parser.parse()
with self.subTest(code=code):
parser = ImportParser(code)
actual = parser.parse()
self.assertEqual(actual, None)
@patch.dict(sys.modules)
@@ -0,0 +1,2 @@
Insert ``import`` after ``from x.y.z <tab>`` in the :term:`REPL`
auto-completion.