Skip to content

Commit c33fefb

Browse files
committed
PyREPL completion: insert "import" after "from foo "
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)
1 parent 3a7df63 commit c33fefb

File tree

3 files changed

+92
-72
lines changed

3 files changed

+92
-72
lines changed

Lib/_pyrepl/_module_completer.py

Lines changed: 36 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ def get_completions(self, line: str) -> tuple[list[str], CompletionAction | None
8888
# no completions are available
8989
return [], None
9090

91-
def complete(self, from_name: str | None, name: str | None) -> tuple[list[str], CompletionAction | None]:
91+
def complete(self, from_name: str | None, name: str | None, space_end: bool) -> tuple[list[str], CompletionAction | None]:
9292
if from_name is None:
9393
# import x.y.z<tab>
9494
assert name is not None
@@ -97,10 +97,17 @@ def complete(self, from_name: str | None, name: str | None) -> tuple[list[str],
9797
return [self.format_completion(path, module) for module in modules], None
9898

9999
if name is None:
100+
if space_end and from_name:
101+
# from x.y.z <tab>
102+
return (["import "], None)
100103
# from x.y.z<tab>
101104
path, prefix = self.get_path_and_prefix(from_name)
102105
modules = self.find_modules(path, prefix)
103-
return [self.format_completion(path, module) for module in modules], None
106+
names = [self.format_completion(path, module) for module in modules]
107+
if len(names) == 1:
108+
# One match: insert a space to allow for "import" suggestion
109+
names[0] = f"{names[0]} "
110+
return names, None
104111

105112
# from x.y import z<tab>
106113
submodules = self.find_modules(from_name, name)
@@ -307,11 +314,12 @@ class ImportParser:
307314
suitable for autocomplete suggestions.
308315
309316
Examples:
310-
- import foo -> Result(from_name=None, name='foo')
311-
- import foo. -> Result(from_name=None, name='foo.')
312-
- from foo -> Result(from_name='foo', name=None)
313-
- from foo import bar -> Result(from_name='foo', name='bar')
314-
- from .foo import ( -> Result(from_name='.foo', name='')
317+
- `import foo` -> Result(from_name=None, name='foo')
318+
- `import foo.` -> Result(from_name=None, name='foo.')
319+
- `from foo` -> Result(from_name='foo', name=None)
320+
- `from foo ` -> Result(from_name='foo', name=None, end_space=True)
321+
- `from foo import bar` -> Result(from_name='foo', name='bar')
322+
- `from .foo import (` -> Result(from_name='.foo', name='')
315323
316324
Note that the parser works in reverse order, starting from the
317325
last token in the input string. This makes the parser more robust
@@ -341,10 +349,10 @@ def __init__(self, code: str) -> None:
341349
tokens = []
342350
self.tokens = TokenQueue(tokens[::-1])
343351

344-
def parse(self) -> tuple[str | None, str | None] | None:
352+
def parse(self) -> tuple[str | None, str | None, bool] | None:
345353
if not (res := self._parse()):
346354
return None
347-
return res.from_name, res.name
355+
return res.from_name, res.name, res.space_end
348356

349357
def _parse(self) -> Result | None:
350358
with self.tokens.save_state():
@@ -354,7 +362,7 @@ def _parse(self) -> Result | None:
354362

355363
def parse_import(self) -> Result:
356364
if self.code.rstrip().endswith('import') and self.code.endswith(' '):
357-
return Result(name='')
365+
return Result(name='', space_end=True)
358366
if self.tokens.peek_string(','):
359367
name = ''
360368
else:
@@ -367,27 +375,32 @@ def parse_import(self) -> Result:
367375
self.tokens.pop()
368376
self.parse_dotted_as_name()
369377
if self.tokens.peek_string('import'):
370-
return Result(name=name)
378+
return Result(name=name, space_end=self.code.endswith(' '))
371379
raise ParseError('parse_import')
372380

373381
def parse_from_import(self) -> Result:
382+
space_end = self.code.endswith(' ')
374383
stripped = self.code.rstrip()
375-
if stripped.endswith('import') and self.code.endswith(' '):
376-
return Result(from_name=self.parse_empty_from_import(), name='')
377-
if stripped.endswith('from') and self.code.endswith(' '):
378-
return Result(from_name='')
384+
if stripped.endswith('import') and space_end:
385+
from_name = self.parse_empty_from_import()
386+
return Result(from_name=from_name, name='', space_end=space_end)
387+
if stripped.endswith('from') and space_end:
388+
return Result(from_name='', space_end=space_end)
379389
if self.tokens.peek_string('(') or self.tokens.peek_string(','):
380-
return Result(from_name=self.parse_empty_from_import(), name='')
381-
if self.code.endswith(' '):
382-
raise ParseError('parse_from_import')
390+
from_name = self.parse_empty_from_import()
391+
return Result(from_name=from_name, name='', space_end=space_end)
383392
name = self.parse_dotted_name()
384393
if '.' in name:
394+
if name.endswith(".") and space_end:
395+
raise ParseError('parse_from_import')
385396
self.tokens.pop_string('from')
386-
return Result(from_name=name)
397+
return Result(from_name=name, space_end=space_end)
387398
if self.tokens.peek_string('from'):
388-
return Result(from_name=name)
399+
return Result(from_name=name, space_end=space_end)
400+
if space_end:
401+
raise ParseError('parse_from_import')
389402
from_name = self.parse_empty_from_import()
390-
return Result(from_name=from_name, name=name)
403+
return Result(from_name=from_name, name=name, space_end=space_end)
391404

392405
def parse_empty_from_import(self) -> str:
393406
if self.tokens.peek_string(','):
@@ -453,10 +466,11 @@ class ParseError(Exception):
453466
pass
454467

455468

456-
@dataclass(frozen=True)
469+
@dataclass(frozen=True, kw_only=True)
457470
class Result:
458471
from_name: str | None = None
459472
name: str | None = None
473+
space_end: bool
460474

461475

462476
class TokenQueue:

Lib/test/test_pyrepl/test_pyrepl.py

Lines changed: 54 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -1168,13 +1168,18 @@ def test_completions(self):
11681168
("import importlib.resources.\t\ta\t\n", "import importlib.resources.abc"),
11691169
("import foo, impo\t\n", "import foo, importlib"),
11701170
("import foo as bar, impo\t\n", "import foo as bar, importlib"),
1171-
("from impo\t\n", "from importlib"),
1172-
("from importlib.res\t\n", "from importlib.resources"),
1173-
("from importlib.\t\tres\t\n", "from importlib.resources"),
1174-
("from importlib.resources.ab\t\n", "from importlib.resources.abc"),
1171+
("from impo\t\n", "from importlib "),
1172+
("from impo\t\t\n", "from importlib import "),
1173+
("from impo \t\n", "from impo import "),
1174+
("from importlib.res\t\n", "from importlib.resources "),
1175+
("from importlib.\t\tres\t\n", "from importlib.resources "),
1176+
("from importlib.res\t\t\n", "from importlib.resources import "),
1177+
("from importlib.res \t\n", "from importlib.res import "),
1178+
("from importlib.resources.ab\t\n", "from importlib.resources.abc "),
11751179
("from importlib import mac\t\n", "from importlib import machinery"),
11761180
("from importlib import res\t\n", "from importlib import resources"),
1177-
("from importlib.res\t import a\t\n", "from importlib.resources import abc"),
1181+
("from importlib.res\timport a\t\n", "from importlib.resources import abc"),
1182+
("from importlib.res\t\ta\t\n", "from importlib.resources import abc"),
11781183
("from __phello__ import s\t\n", "from __phello__ import spam"), # frozen module
11791184
)
11801185
for code, expected in cases:
@@ -1191,10 +1196,10 @@ def test_private_completions(self):
11911196
cases = (
11921197
# Return public methods by default
11931198
("import \t\n", "import public"),
1194-
("from \t\n", "from public"),
1199+
("from \t\n", "from public "),
11951200
# Return private methods if explicitly specified
11961201
("import _\t\n", "import _private"),
1197-
("from _\t\n", "from _private"),
1202+
("from _\t\n", "from _private "),
11981203
)
11991204
for code, expected in cases:
12001205
with self.subTest(code=code):
@@ -1227,7 +1232,7 @@ def test_sub_module_private_completions(self):
12271232
def test_builtin_completion_top_level(self):
12281233
cases = (
12291234
("import bui\t\n", "import builtins"),
1230-
("from bui\t\n", "from builtins"),
1235+
("from bui\t\n", "from builtins "),
12311236
)
12321237
for code, expected in cases:
12331238
with self.subTest(code=code):
@@ -1240,11 +1245,11 @@ def test_relative_completions(self):
12401245
cases = (
12411246
(None, "from .readl\t\n", "from .readl"),
12421247
(None, "from . import readl\t\n", "from . import readl"),
1243-
("_pyrepl", "from .readl\t\n", "from .readline"),
1248+
("_pyrepl", "from .readl\t\n", "from .readline "),
12441249
("_pyrepl", "from . import readl\t\n", "from . import readline"),
12451250
("_pyrepl", "from .readline import mul\t\n", "from .readline import multiline_input"),
12461251
("_pyrepl", "from .. import toodeep\t\n", "from .. import toodeep"),
1247-
("concurrent", "from .futures.i\t\n", "from .futures.interpreter"),
1252+
("concurrent", "from .futures.i\t\n", "from .futures.interpreter "),
12481253
)
12491254
for package, code, expected in cases:
12501255
with self.subTest(code=code):
@@ -1545,41 +1550,43 @@ def test_get_path_and_prefix(self):
15451550

15461551
def test_parse(self):
15471552
cases = (
1548-
('import ', (None, '')),
1549-
('import foo', (None, 'foo')),
1550-
('import foo,', (None, '')),
1551-
('import foo, ', (None, '')),
1552-
('import foo, bar', (None, 'bar')),
1553-
('import foo, bar, baz', (None, 'baz')),
1554-
('import foo as bar,', (None, '')),
1555-
('import foo as bar, ', (None, '')),
1556-
('import foo as bar, baz', (None, 'baz')),
1557-
('import a.', (None, 'a.')),
1558-
('import a.b', (None, 'a.b')),
1559-
('import a.b.', (None, 'a.b.')),
1560-
('import a.b.c', (None, 'a.b.c')),
1561-
('import a.b.c, foo', (None, 'foo')),
1562-
('import a.b.c, foo.bar', (None, 'foo.bar')),
1563-
('import a.b.c, foo.bar,', (None, '')),
1564-
('import a.b.c, foo.bar, ', (None, '')),
1565-
('from foo', ('foo', None)),
1566-
('from a.', ('a.', None)),
1567-
('from a.b', ('a.b', None)),
1568-
('from a.b.', ('a.b.', None)),
1569-
('from a.b.c', ('a.b.c', None)),
1570-
('from foo import ', ('foo', '')),
1571-
('from foo import a', ('foo', 'a')),
1572-
('from ', ('', None)),
1573-
('from . import a', ('.', 'a')),
1574-
('from .foo import a', ('.foo', 'a')),
1575-
('from ..foo import a', ('..foo', 'a')),
1576-
('from foo import (', ('foo', '')),
1577-
('from foo import ( ', ('foo', '')),
1578-
('from foo import (a', ('foo', 'a')),
1579-
('from foo import (a,', ('foo', '')),
1580-
('from foo import (a, ', ('foo', '')),
1581-
('from foo import (a, c', ('foo', 'c')),
1582-
('from foo import (a as b, c', ('foo', 'c')),
1553+
('import ', (None, '', True)),
1554+
('import foo', (None, 'foo', False)),
1555+
('import foo,', (None, '', False)),
1556+
('import foo, ', (None, '', True)),
1557+
('import foo, bar', (None, 'bar', False)),
1558+
('import foo, bar, baz', (None, 'baz', False)),
1559+
('import foo as bar,', (None, '', False)),
1560+
('import foo as bar, ', (None, '', True)),
1561+
('import foo as bar, baz', (None, 'baz', False)),
1562+
('import a.', (None, 'a.', False)),
1563+
('import a.b', (None, 'a.b', False)),
1564+
('import a.b.', (None, 'a.b.', False)),
1565+
('import a.b.c', (None, 'a.b.c', False)),
1566+
('import a.b.c, foo', (None, 'foo', False)),
1567+
('import a.b.c, foo.bar', (None, 'foo.bar', False)),
1568+
('import a.b.c, foo.bar,', (None, '', False)),
1569+
('import a.b.c, foo.bar, ', (None, '', True)),
1570+
('from foo', ('foo', None, False)),
1571+
('from foo ', ('foo', None, True)),
1572+
('from a.', ('a.', None, False)),
1573+
('from a.b', ('a.b', None, False)),
1574+
('from a.b.', ('a.b.', None, False)),
1575+
('from a.b ', ('a.b', None, True)),
1576+
('from a.b.c', ('a.b.c', None, False)),
1577+
('from foo import ', ('foo', '', True)),
1578+
('from foo import a', ('foo', 'a', False)),
1579+
('from ', ('', None, True)),
1580+
('from . import a', ('.', 'a', False)),
1581+
('from .foo import a', ('.foo', 'a', False)),
1582+
('from ..foo import a', ('..foo', 'a', False)),
1583+
('from foo import (', ('foo', '', False)),
1584+
('from foo import ( ', ('foo', '', True)),
1585+
('from foo import (a', ('foo', 'a', False)),
1586+
('from foo import (a,', ('foo', '', False)),
1587+
('from foo import (a, ', ('foo', '', True)),
1588+
('from foo import (a, c', ('foo', 'c', False)),
1589+
('from foo import (a as b, c', ('foo', 'c', False)),
15831590
)
15841591
for code, parsed in cases:
15851592
parser = ImportParser(code)
@@ -1603,12 +1610,9 @@ def test_parse_error(self):
16031610
cases = (
16041611
'',
16051612
'import foo ',
1606-
'from foo ',
16071613
'import foo. ',
16081614
'import foo.bar ',
1609-
'from foo ',
16101615
'from foo. ',
1611-
'from foo.bar ',
16121616
'from foo import bar ',
16131617
'from foo import (bar ',
16141618
'from foo import bar, baz ',
@@ -1647,9 +1651,9 @@ def test_parse_error(self):
16471651
'if 1:\n pass\n\tpass', # _tokenize TabError -> tokenize TabError
16481652
)
16491653
for code in cases:
1650-
parser = ImportParser(code)
1651-
actual = parser.parse()
16521654
with self.subTest(code=code):
1655+
parser = ImportParser(code)
1656+
actual = parser.parse()
16531657
self.assertEqual(actual, None)
16541658

16551659
@patch.dict(sys.modules)
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Insert ``import`` after ``from x.y.z <tab>`` in the :term:`REPL`
2+
auto-completion.

0 commit comments

Comments
 (0)