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>

Author: loic-simon

Lib/_pyrepl/_module_completer.py

@@ -97,6 +97,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]:
@@ -111,10 +112,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 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 @@ 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():
@@ -414,7 +422,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:
@@ -427,27 +435,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(','):
@@ -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:

Lib/test/test_pyrepl/test_pyrepl.py

@@ -1231,12 +1231,20 @@ def test_completions(self):
             ("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 @@ 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 .', ('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 @@ 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 ',
@@ -1723,9 +1730,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.