gh-140870: PyREPL auto-complete module attributes in import statements (#140871)

Co-authored-by: Pablo Galindo Salgado <pablogsal@gmail.com>

Author: loic-simon

Lib/_pyrepl/_module_completer.py

@@ -3,6 +3,7 @@
 import importlib
 import os
 import pkgutil
+import re
 import sys
 import token
 import tokenize
@@ -16,7 +17,9 @@
 TYPE_CHECKING = False
 
 if TYPE_CHECKING:
+    from types import ModuleType
     from typing import Any, Iterable, Iterator, Mapping
+    from .types import CompletionAction
 
 
 HARDCODED_SUBMODULES = {
@@ -28,6 +31,17 @@
     "xml.parsers.expat": ["errors", "model"],
 }
 
+AUTO_IMPORT_DENYLIST = {
+    # Standard library modules/submodules that have import side effects
+    # and must not be automatically imported to complete attributes
+    re.compile(r"antigravity"),  # Calls webbrowser.open
+    re.compile(r"idlelib\..+"),  # May open IDLE GUI
+    re.compile(r"test\..+"),  # Various side-effects
+    re.compile(r"this"),  # Prints to stdout
+    re.compile(r"_ios_support"),  # Spawns a subprocess
+    re.compile(r".+\.__main__"),  # Should not be imported
+}
+
 
 def make_default_module_completer() -> ModuleCompleter:
     # Inside pyrepl, __package__ is set to None by default
@@ -53,11 +67,17 @@ class ModuleCompleter:
     def __init__(self, namespace: Mapping[str, Any] | None = None) -> None:
         self.namespace = namespace or {}
         self._global_cache: list[pkgutil.ModuleInfo] = []
+        self._failed_imports: set[str] = set()
         self._curr_sys_path: list[str] = sys.path[:]
         self._stdlib_path = os.path.dirname(importlib.__path__[0])
 
-    def get_completions(self, line: str) -> list[str] | None:
-        """Return the next possible import completions for 'line'."""
+    def get_completions(self, line: str) -> tuple[list[str], CompletionAction | None] | None:
+        """Return the next possible import completions for 'line'.
+
+        For attributes completion, if the module to complete from is not
+        imported, also return an action (prompt + callback to run if the
+        user press TAB again) to import the module.
+        """
         result = ImportParser(line).parse()
         if not result:
             return None
@@ -66,24 +86,26 @@ def get_completions(self, line: str) -> list[str] | None:
         except Exception:
             # Some unexpected error occurred, make it look like
             # no completions are available
-            return []
+            return [], None
 
-    def complete(self, from_name: str | None, name: str | None) -> list[str]:
+    def complete(self, from_name: str | None, name: str | None) -> tuple[list[str], CompletionAction | None]:
         if from_name is None:
             # import x.y.z<tab>
             assert name is not None
             path, prefix = self.get_path_and_prefix(name)
             modules = self.find_modules(path, prefix)
-            return [self.format_completion(path, module) for module in modules]
+            return [self.format_completion(path, module) for module in modules], None
 
         if name is None:
             # from x.y.z<tab>
             path, prefix = self.get_path_and_prefix(from_name)
             modules = self.find_modules(path, prefix)
-            return [self.format_completion(path, module) for module in modules]
+            return [self.format_completion(path, module) for module in modules], None
 
         # from x.y import z<tab>
-        return self.find_modules(from_name, name)
+        submodules = self.find_modules(from_name, name)
+        attributes, action = self.find_attributes(from_name, name)
+        return sorted({*submodules, *attributes}), action
 
     def find_modules(self, path: str, prefix: str) -> list[str]:
         """Find all modules under 'path' that start with 'prefix'."""
@@ -101,23 +123,25 @@ def _find_modules(self, path: str, prefix: str) -> list[str]:
                                    if self.is_suggestion_match(module.name, prefix)]
             return sorted(builtin_modules + third_party_modules)
 
-        if path.startswith('.'):
-            # Convert relative path to absolute path
-            package = self.namespace.get('__package__', '')
-            path = self.resolve_relative_name(path, package)  # type: ignore[assignment]
-            if path is None:
-                return []
+        path = self._resolve_relative_path(path)  # type: ignore[assignment]
+        if path is None:
+            return []
 
         modules: Iterable[pkgutil.ModuleInfo] = self.global_cache
         imported_module = sys.modules.get(path.split('.')[0])
         if imported_module:
-            # Filter modules to those who name and specs match the
+            # Filter modules to those whose name and specs match the
             # imported module to avoid invalid suggestions
             spec = imported_module.__spec__
             if spec:
+                def _safe_find_spec(mod: pkgutil.ModuleInfo) -> bool:
+                    try:
+                        return mod.module_finder.find_spec(mod.name, None) == spec
+                    except Exception:
+                        return False
                 modules = [mod for mod in modules
                            if mod.name == spec.name
-                           and mod.module_finder.find_spec(mod.name, None) == spec]
+                           and _safe_find_spec(mod)]
             else:
                 modules = []
 
@@ -142,6 +166,32 @@ def _is_stdlib_module(self, module_info: pkgutil.ModuleInfo) -> bool:
         return (isinstance(module_info.module_finder, FileFinder)
                 and module_info.module_finder.path == self._stdlib_path)
 
+    def find_attributes(self, path: str, prefix: str) -> tuple[list[str], CompletionAction | None]:
+        """Find all attributes of module 'path' that start with 'prefix'."""
+        attributes, action = self._find_attributes(path, prefix)
+        # Filter out invalid attribute names
+        # (for example those containing dashes that cannot be imported with 'import')
+        return [attr for attr in attributes if attr.isidentifier()], action
+
+    def _find_attributes(self, path: str, prefix: str) -> tuple[list[str], CompletionAction | None]:
+        path = self._resolve_relative_path(path)  # type: ignore[assignment]
+        if path is None:
+            return [], None
+
+        imported_module = sys.modules.get(path)
+        if not imported_module:
+            if path in self._failed_imports:  # Do not propose to import again
+                return [], None
+            imported_module = self._maybe_import_module(path)
+        if not imported_module:
+            return [], self._get_import_completion_action(path)
+        try:
+            module_attributes = dir(imported_module)
+        except Exception:
+            module_attributes = []
+        return [attr_name for attr_name in module_attributes
+                if self.is_suggestion_match(attr_name, prefix)], None
+
     def is_suggestion_match(self, module_name: str, prefix: str) -> bool:
         if prefix:
             return module_name.startswith(prefix)
@@ -186,6 +236,13 @@ def format_completion(self, path: str, module: str) -> str:
             return f'{path}{module}'
         return f'{path}.{module}'
 
+    def _resolve_relative_path(self, path: str) -> str | None:
+        """Resolve a relative import path to absolute. Returns None if unresolvable."""
+        if path.startswith('.'):
+            package = self.namespace.get('__package__', '')
+            return self.resolve_relative_name(path, package)
+        return path
+
     def resolve_relative_name(self, name: str, package: str) -> str | None:
         """Resolve a relative module name to an absolute name.
 
@@ -210,8 +267,39 @@ def global_cache(self) -> list[pkgutil.ModuleInfo]:
         if not self._global_cache or self._curr_sys_path != sys.path:
             self._curr_sys_path = sys.path[:]
             self._global_cache = list(pkgutil.iter_modules())
+            self._failed_imports.clear()  # retry on sys.path change
         return self._global_cache
 
+    def _maybe_import_module(self, fqname: str) -> ModuleType | None:
+        if any(pattern.fullmatch(fqname) for pattern in AUTO_IMPORT_DENYLIST):
+            # Special-cased modules with known import side-effects
+            return None
+        root = fqname.split(".")[0]
+        mod_info = next((m for m in self.global_cache if m.name == root), None)
+        if not mod_info or not self._is_stdlib_module(mod_info):
+            # Only import stdlib modules (no risk of import side-effects)
+            return None
+        try:
+            return importlib.import_module(fqname)
+        except Exception:
+            sys.modules.pop(fqname, None)  # Clean half-imported module
+            return None
+
+    def _get_import_completion_action(self, path: str) -> CompletionAction:
+        prompt = ("[ module not imported, press again to import it "
+                  "and propose attributes ]")
+
+        def _do_import() -> str | None:
+            try:
+                importlib.import_module(path)
+                return None
+            except Exception as exc:
+                sys.modules.pop(path, None)  # Clean half-imported module
+                self._failed_imports.add(path)
+                return f"[ error during import: {exc} ]"
+
+        return (prompt, _do_import)
+
 
 class ImportParser:
     """

Lib/_pyrepl/completing_reader.py

@@ -29,8 +29,9 @@
 
 # types
 Command = commands.Command
-if False:
-    from .types import KeySpec, CommandName
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from .types import KeySpec, CommandName, CompletionAction
 
 
 def prefix(wordlist: list[str], j: int = 0) -> str:
@@ -168,15 +169,25 @@ def do(self) -> None:
         r: CompletingReader
         r = self.reader  # type: ignore[assignment]
         last_is_completer = r.last_command_is(self.__class__)
+        if r.cmpltn_action:
+            if last_is_completer:  # double-tab: execute action
+                msg = r.cmpltn_action[1]()
+                r.cmpltn_action = None  # consumed
+                if msg:
+                    r.msg = msg
+            else:  # other input since last tab: cancel action
+                r.cmpltn_action = None
+
         immutable_completions = r.assume_immutable_completions
         completions_unchangable = last_is_completer and immutable_completions
         stem = r.get_stem()
         if not completions_unchangable:
-            r.cmpltn_menu_choices = r.get_completions(stem)
+            r.cmpltn_menu_choices, r.cmpltn_action = r.get_completions(stem)
 
         completions = r.cmpltn_menu_choices
         if not completions:
-            r.error("no matches")
+            if not r.cmpltn_action:
+                r.error("no matches")
         elif len(completions) == 1:
             completion = stripcolor(completions[0])
             if completions_unchangable and len(completion) == len(stem):
@@ -204,6 +215,16 @@ def do(self) -> None:
                     r.msg = "[ not unique ]"
                     r.dirty = True
 
+        if r.cmpltn_action:
+            if r.msg and r.cmpltn_message_visible:
+                # There is already a message (eg. [ not unique ]) that
+                # would conflict for next tab: cancel action
+                r.cmpltn_action = None
+            else:
+                r.msg = r.cmpltn_action[0]
+                r.cmpltn_message_visible = True
+                r.dirty = True
+
 
 class self_insert(commands.self_insert):
     def do(self) -> None:
@@ -242,6 +263,7 @@ class CompletingReader(Reader):
     cmpltn_message_visible: bool = field(init=False)
     cmpltn_menu_end: int = field(init=False)
     cmpltn_menu_choices: list[str] = field(init=False)
+    cmpltn_action: CompletionAction | None = field(init=False)
 
     def __post_init__(self) -> None:
         super().__post_init__()
@@ -283,6 +305,7 @@ def cmpltn_reset(self) -> None:
         self.cmpltn_message_visible = False
         self.cmpltn_menu_end = 0
         self.cmpltn_menu_choices = []
+        self.cmpltn_action = None
 
     def get_stem(self) -> str:
         st = self.syntax_table
@@ -293,8 +316,8 @@ def get_stem(self) -> str:
             p -= 1
         return ''.join(b[p+1:self.pos])
 
-    def get_completions(self, stem: str) -> list[str]:
-        return []
+    def get_completions(self, stem: str) -> tuple[list[str], CompletionAction | None]:
+        return [], None
 
     def get_line(self) -> str:
         """Return the current line until the cursor position."""

Lib/_pyrepl/reader.py

@@ -381,9 +381,17 @@ def calc_screen(self) -> list[str]:
         self.screeninfo = screeninfo
         self.cxy = self.pos2xy()
         if self.msg:
+            width = self.console.width
             for mline in self.msg.split("\n"):
-                screen.append(mline)
-                screeninfo.append((0, []))
+                # If self.msg is larger than console width, make it fit
+                # TODO: try to split between words?
+                if not mline:
+                    screen.append("")
+                    screeninfo.append((0, []))
+                    continue
+                for r in range((len(mline) - 1) // width + 1):
+                    screen.append(mline[r * width : (r + 1) * width])
+                    screeninfo.append((0, []))
 
         self.last_refresh_cache.update_cache(self, screen, screeninfo)
         return screen
@@ -628,7 +636,6 @@ def suspend_colorization(self) -> SimpleContextManager:
         finally:
             self.can_colorize = old_can_colorize
 
-
     def finish(self) -> None:
         """Called when a command signals that we're finished."""
         pass

Lib/_pyrepl/readline.py

@@ -56,7 +56,7 @@
 # types
 Command = commands.Command
 from collections.abc import Callable, Collection
-from .types import Callback, Completer, KeySpec, CommandName
+from .types import Callback, Completer, KeySpec, CommandName, CompletionAction
 
 TYPE_CHECKING = False
 
@@ -135,7 +135,7 @@ def get_stem(self) -> str:
             p -= 1
         return "".join(b[p + 1 : self.pos])
 
-    def get_completions(self, stem: str) -> list[str]:
+    def get_completions(self, stem: str) -> tuple[list[str], CompletionAction | None]:
         module_completions = self.get_module_completions()
         if module_completions is not None:
             return module_completions
@@ -145,7 +145,7 @@ def get_completions(self, stem: str) -> list[str]:
             while p > 0 and b[p - 1] != "\n":
                 p -= 1
             num_spaces = 4 - ((self.pos - p) % 4)
-            return [" " * num_spaces]
+            return [" " * num_spaces], None
         result = []
         function = self.config.readline_completer
         if function is not None:
@@ -166,9 +166,9 @@ def get_completions(self, stem: str) -> list[str]:
             # emulate the behavior of the standard readline that sorts
             # the completions before displaying them.
             result.sort()
-        return result
+        return result, None
 
-    def get_module_completions(self) -> list[str] | None:
+    def get_module_completions(self) -> tuple[list[str], CompletionAction | None] | None:
         line = self.get_line()
         return self.config.module_completer.get_completions(line)
 

Lib/_pyrepl/types.py

@@ -8,3 +8,4 @@
 type Completer = Callable[[str, int], str | None]
 type CharBuffer = list[str]
 type CharWidths = list[int]
+type CompletionAction = tuple[str, Callable[[], str | None]]