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)

Author: loic-simon

Lib/_pyrepl/_module_completer.py

@@ -93,6 +93,7 @@ def complete(
         self,
         from_name: str | None,
         name: str | None,
+        space_end: bool,
         *,
         include_values: bool = True,
     ) -> tuple[list[str], list[Any], CompletionAction | None]:
@@ -107,10 +108,16 @@ def complete(
             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 len(names) == 1:
+                # One match: insert a space to allow for "import" suggestion
+                names[0] = f"{names[0]} "
             # These are always modules, use dummy values to get the right color
             values = [sys] * len(names) if include_values else []
             return names, values, None
@@ -343,11 +350,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, end_space=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
@@ -377,10 +385,10 @@ def __init__(self, code: str) -> None:
             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():
@@ -390,7 +398,7 @@ def _parse(self) -> Result | None:
 
     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:
@@ -403,27 +411,32 @@ def parse_import(self) -> Result:
             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(','):
@@ -489,10 +502,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:

Lib/test/test_pyrepl/test_pyrepl.py

@@ -1230,13 +1230,18 @@ def test_completions(self):
             ("import importlib.resources.\t\ta\t\n", "import importlib.resources.abc"),
             ("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 importlib.res\t\n", "from importlib.resources"),
-            ("from importlib.\t\tres\t\n", "from importlib.resources"),
-            ("from importlib.resources.ab\t\n", "from importlib.resources.abc"),
+            ("from impo\t\n", "from importlib "),
+            ("from impo\t\t\n", "from importlib import "),
+            ("from impo \t\n", "from impo 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.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\timport 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
         )
         for code, expected in cases:
@@ -1253,10 +1258,10 @@ def test_private_completions(self):
         cases = (
             # Return public methods by default
             ("import \t\n", "import public"),
-            ("from \t\n", "from public"),
+            ("from \t\n", "from public "),
             # Return private methods if explicitly specified
             ("import _\t\n", "import _private"),
-            ("from _\t\n", "from _private"),
+            ("from _\t\n", "from _private "),
         )
         for code, expected in cases:
             with self.subTest(code=code):
@@ -1289,7 +1294,7 @@ def test_sub_module_private_completions(self):
     def test_builtin_completion_top_level(self):
         cases = (
             ("import bui\t\n", "import builtins"),
-            ("from bui\t\n", "from builtins"),
+            ("from bui\t\n", "from builtins "),
         )
         for code, expected in cases:
             with self.subTest(code=code):
@@ -1302,11 +1307,11 @@ def test_relative_completions(self):
         cases = (
             (None, "from .readl\t\n", "from .readl"),
             (None, "from . import readl\t\n", "from . import readl"),
-            ("_pyrepl", "from .readl\t\n", "from .readline"),
+            ("_pyrepl", "from .readl\t\n", "from .readline "),
             ("_pyrepl", "from . import readl\t\n", "from . import readline"),
             ("_pyrepl", "from .readline import mul\t\n", "from .readline import multiline_input"),
             ("_pyrepl", "from .. import toodeep\t\n", "from .. import toodeep"),
-            ("concurrent", "from .futures.i\t\n", "from .futures.interpreter"),
+            ("concurrent", "from .futures.i\t\n", "from .futures.interpreter "),
         )
         for package, code, expected in cases:
             with self.subTest(code=code):
@@ -1610,41 +1615,43 @@ def test_get_path_and_prefix(self):
 
     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.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.b', ('a.b', None, False)),
+            ('from a.b.', ('a.b.', None, False)),
+            ('from a.b ', ('a.b', None, True)),
+            ('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)
@@ -1668,12 +1675,9 @@ def test_parse_error(self):
         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 ',
@@ -1712,9 +1716,9 @@ def test_parse_error(self):
             '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)

Misc/NEWS.d/next/Core_and_Builtins/2026-04-12-13-48-26.gh-issue-69605.T9HWOL.rst

@@ -0,0 +1,2 @@
+Insert ``import`` after ``from x.y.z <tab>`` in the :term:`REPL`
+auto-completion.