gh-130472: Use fancycompleter in import completions

Lib/_pyrepl/_module_completer.py

@@ -13,6 +13,7 @@
 from dataclasses import dataclass
 from itertools import chain
 from tokenize import TokenInfo
+from .fancycompleter import safe_getattr
 
 TYPE_CHECKING = False
 
@@ -71,7 +72,7 @@ def __init__(self, namespace: Mapping[str, Any] | None = None) -> None:
         self._curr_sys_path: list[str] = sys.path[:]
         self._stdlib_path = os.path.dirname(importlib.__path__[0])
 
-    def get_completions(self, line: str) -> tuple[list[str], CompletionAction | None] | None:
+    def get_completions(self, line: str) -> tuple[list[str], list[Any], CompletionAction | None] | None:
         """Return the next possible import completions for 'line'.
 
         For attributes completion, if the module to complete from is not
@@ -86,26 +87,40 @@ def get_completions(self, line: str) -> tuple[list[str], CompletionAction | None
         except Exception:
             # Some unexpected error occurred, make it look like
             # no completions are available
-            return [], None
+            return [], [], None
 
-    def complete(self, from_name: str | None, name: str | None) -> tuple[list[str], CompletionAction | None]:
+    def complete(self, from_name: str | None, name: str | None) -> tuple[list[str], list[Any], 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], None
+            names = [self.format_completion(path, module) for module in modules]
+            # These are always modules, use dummy values to get the right color
+            values = [sys] * len(names)
+            return names, values, 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], None
+            names = [self.format_completion(path, module) for module in modules]
+            # These are always modules, use dummy values to get the right color
+            values = [sys] * len(names)
+            return names, values, None
 
         # from x.y import z<tab>
         submodules = self.find_modules(from_name, name)
-        attributes, action = self.find_attributes(from_name, name)
-        return sorted({*submodules, *attributes}), action
+        attr_names, attr_values, action = self.find_attributes(from_name, name)
+        all_names = sorted({*submodules, *attr_names})
+        # Build values list matching the sorted order:
+        # submodules use `sys` as a dummy value so they get the 'module' color,
+        # attributes use their actual value.
+        submodule_set = set(submodules)
+        attr_map = dict(zip(attr_names, attr_values))
+        all_values = [attr_map.get(n) if n not in submodule_set else sys
+                      for n in all_names]
+        return all_names, all_values, action
 
     def find_modules(self, path: str, prefix: str) -> list[str]:
         """Find all modules under 'path' that start with 'prefix'."""
@@ -166,31 +181,43 @@ 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]:
+    def find_attributes(self, path: str, prefix: str) -> tuple[list[str], list[Any], CompletionAction | None]:
         """Find all attributes of module 'path' that start with 'prefix'."""
-        attributes, action = self._find_attributes(path, prefix)
+        attributes, values, 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]:
+        filtered_names = []
+        filtered_values = []
+        for attr, val in zip(attributes, values):
+            if attr.isidentifier():
+                filtered_names.append(attr)
+                filtered_values.append(val)
+        return filtered_names, filtered_values, action
+
+    def _find_attributes(self, path: str, prefix: str) -> tuple[list[str], list[Any], CompletionAction | None]:
         path = self._resolve_relative_path(path)  # type: ignore[assignment]
         if path is None:
-            return [], 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
+                return [], [], None
             imported_module = self._maybe_import_module(path)
         if not imported_module:
-            return [], self._get_import_completion_action(path)
+            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
+        names = []
+        values = []
+        for attr_name in module_attributes:
+            if not self.is_suggestion_match(attr_name, prefix):
+                continue
+            names.append(attr_name)
+            values.append(safe_getattr(imported_module, attr_name))
+        return names, values, None
 
     def is_suggestion_match(self, module_name: str, prefix: str) -> bool:
         if prefix:

Lib/_pyrepl/fancycompleter.py

@@ -8,6 +8,40 @@
 import keyword
 import types
 
+
+def safe_getattr(obj, name):
+    # Mirror rlcompleter's safeguards so completion does not
+    # call properties or reify lazy module attributes.
+    if isinstance(getattr(type(obj), name, None), property):
+        return None
+    if (isinstance(obj, types.ModuleType)
+        and isinstance(obj.__dict__.get(name), types.LazyImportType)
+    ):
+        return obj.__dict__.get(name)
+    return getattr(obj, name, None)
+
+
+def colorize_matches(names, values, theme):
+    return [
+        _color_for_obj(name, obj, theme)
+        for name, obj in zip(names, values)
+    ]
+
+def _color_for_obj(name, value, theme):
+    t = type(value)
+    color = _color_by_type(t, theme)
+    return f"{color}{name}{ANSIColors.RESET}"
+
+
+def _color_by_type(t, theme):
+    typename = t.__name__
+    # this is needed e.g. to turn method-wrapper into method_wrapper,
+    # because if we want _colorize.FancyCompleter to be "dataclassable"
+    # our keys need to be valid identifiers.
+    typename = typename.replace('-', '_').replace('.', '_')
+    return getattr(theme.fancycompleter, typename, ANSIColors.RESET)
+
+
 class Completer(rlcompleter.Completer):
     """
     When doing something like a.b.<tab>, keep the full a.b.attr completion
@@ -143,21 +177,7 @@ def _attr_matches(self, text):
                     word[:n] == attr
                     and not (noprefix and word[:n+1] == noprefix)
                 ):
-                    # Mirror rlcompleter's safeguards so completion does not
-                    # call properties or reify lazy module attributes.
-                    if isinstance(getattr(type(thisobject), word, None), property):
-                        value = None
-                    elif (
-                        isinstance(thisobject, types.ModuleType)
-                        and isinstance(
-                            thisobject.__dict__.get(word),
-                            types.LazyImportType,
-                        )
-                    ):
-                        value = thisobject.__dict__.get(word)
-                    else:
-                        value = getattr(thisobject, word, None)
-
+                    value = safe_getattr(thisobject, word)
                     names.append(word)
                     values.append(value)
             if names or not noprefix:
@@ -170,23 +190,7 @@ def _attr_matches(self, text):
         return expr, attr, names, values
 
     def colorize_matches(self, names, values):
-        return [
-            self._color_for_obj(name, obj)
-            for name, obj in zip(names, values)
-        ]
-
-    def _color_for_obj(self, name, value):
-        t = type(value)
-        color = self._color_by_type(t)
-        return f"{color}{name}{ANSIColors.RESET}"
-
-    def _color_by_type(self, t):
-        typename = t.__name__
-        # this is needed e.g. to turn method-wrapper into method_wrapper,
-        # because if we want _colorize.FancyCompleter to be "dataclassable"
-        # our keys need to be valid identifiers.
-        typename = typename.replace('-', '_').replace('.', '_')
-        return getattr(self.theme.fancycompleter, typename, ANSIColors.RESET)
+        return colorize_matches(names, values, self.theme)
 
 
 def commonprefix(names):

Lib/_pyrepl/readline.py

@@ -40,7 +40,7 @@
 from .completing_reader import CompletingReader, stripcolor
 from .console import Console as ConsoleType
 from ._module_completer import ModuleCompleter, make_default_module_completer
-from .fancycompleter import Completer as FancyCompleter
+from .fancycompleter import Completer as FancyCompleter, colorize_matches
 
 Console: type[ConsoleType]
 _error: tuple[type[Exception], ...] | type[Exception]
@@ -104,6 +104,7 @@ class ReadlineConfig:
     readline_completer: Completer | None = None
     completer_delims: frozenset[str] = frozenset(" \t\n`~!@#$%^&*()-=+[{]}\\|;:'\",<>/?")
     module_completer: ModuleCompleter = field(default_factory=make_default_module_completer)
+    colorize_completions: Callable[[list[str], list[object]], list[str]] | None = None
 
 @dataclass(kw_only=True)
 class ReadlineAlikeReader(historical_reader.HistoricalReader, CompletingReader):
@@ -169,8 +170,14 @@ def get_completions(self, stem: str) -> tuple[list[str], CompletionAction | None
         return result, None
 
     def get_module_completions(self) -> tuple[list[str], CompletionAction | None] | None:
-        line = self.get_line()
-        return self.config.module_completer.get_completions(line)
+        line = stripcolor(self.get_line())
+        result = self.config.module_completer.get_completions(line)
+        if result is None:
+            return None
+        names, values, action = result
+        if self.config.colorize_completions:
+            names = self.config.colorize_completions(names, values)
+        return names, action
 
     def get_trimmed_history(self, maxlength: int) -> list[str]:
         if maxlength >= 0:
@@ -609,13 +616,18 @@ def _setup(namespace: Mapping[str, Any]) -> None:
     # set up namespace in rlcompleter, which requires it to be a bona fide dict
     if not isinstance(namespace, dict):
         namespace = dict(namespace)
-    _wrapper.config.module_completer = ModuleCompleter(namespace)
     use_basic_completer = (
         not sys.flags.ignore_environment
         and os.getenv("PYTHON_BASIC_COMPLETER")
     )
     completer_cls = RLCompleter if use_basic_completer else FancyCompleter
-    _wrapper.config.readline_completer = completer_cls(namespace).complete
+    completer = completer_cls(namespace)
+    _wrapper.config.readline_completer = completer.complete
+    if getattr(completer, 'use_colors', False):
+        def _colorize(names: list[str], values: list[object]) -> list[str]:
+            return colorize_matches(names, values, completer.theme)
+        _wrapper.config.colorize_completions = _colorize
+    _wrapper.config.module_completer = ModuleCompleter(namespace)
 
     # this is not really what readline.c does.  Better than nothing I guess
     import builtins

Lib/test/test_pyrepl/test_fancycompleter.py

@@ -5,7 +5,7 @@
 
 from _colorize import ANSIColors, get_theme
 from _pyrepl.completing_reader import stripcolor
-from _pyrepl.fancycompleter import Completer, commonprefix
+from _pyrepl.fancycompleter import Completer, commonprefix, _color_for_obj
 from test.support.import_helper import ready_to_import
 
 class MockPatch:
@@ -168,8 +168,8 @@ def test_complete_global_colored(self):
         self.assertEqual(compl.global_matches('nothing'), [])
 
     def test_colorized_match_is_stripped(self):
-        compl = Completer({'a': 42}, use_colors=True)
-        match = compl._color_for_obj('spam', 1)
+        theme = get_theme()
+        match = _color_for_obj('spam', 1, theme)
         self.assertEqual(stripcolor(match), 'spam')
 
     def test_complete_with_indexer(self):

Lib/test/test_pyrepl/test_pyrepl.py

@@ -35,14 +35,15 @@
     multiline_input,
     code_to_events,
 )
+from _colorize import ANSIColors, get_theme
 from _pyrepl.console import Event
 from _pyrepl.completing_reader import stripcolor
 from _pyrepl._module_completer import (
     ImportParser,
     ModuleCompleter,
     HARDCODED_SUBMODULES,
 )
-from _pyrepl.fancycompleter import Completer as FancyCompleter
+from _pyrepl.fancycompleter import Completer as FancyCompleter, colorize_matches
 import _pyrepl.readline as pyrepl_readline
 from _pyrepl.readline import (
     ReadlineAlikeReader,
@@ -1629,7 +1630,7 @@ def test_suggestions_and_messages(self) -> None:
                         result = completer.get_completions(code)
                         self.assertEqual(result is None, expected is None)
                         if result:
-                            compl, act = result
+                            compl, _values, act = result
                             self.assertEqual(compl, expected[0])
                             self.assertEqual(act is None, expected[1] is None)
                             if act:
@@ -1641,6 +1642,39 @@ def test_suggestions_and_messages(self) -> None:
                         new_imports = sys.modules.keys() - _imported
                         self.assertSetEqual(new_imports, expected_imports)
 
+    def test_colorize_import_completions(self) -> None:
+        theme = get_theme()
+        type_color = theme.fancycompleter.type
+        module_color = theme.fancycompleter.module
+        R = ANSIColors.RESET
+
+        colorize = lambda names, values: colorize_matches(names, values, theme)
+        config = ReadlineConfig(colorize_completions=colorize)
+        reader = ReadlineAlikeReader(
+            console=FakeConsole(events=[]),
+            config=config,
+        )
+
+        # "from collections import de" -> defaultdict (type) and deque (type)
+        reader.buffer = list("from collections import de")
+        reader.pos = len(reader.buffer)
+        names, action = reader.get_module_completions()
+        self.assertEqual(names, [
+            f"{type_color}defaultdict{R}",
+            f"{type_color}deque{R}",
+        ])
+        self.assertIsNone(action)
+
+        # "from importlib.m" has submodule completions colored as modules
+        reader.buffer = list("from importlib.m")
+        reader.pos = len(reader.buffer)
+        names, action = reader.get_module_completions()
+        self.assertEqual(names, [
+            f"{module_color}importlib.machinery{R}",
+            f"{module_color}importlib.metadata{R}",
+        ])
+        self.assertIsNone(action)
+
 
 # Audit hook used to check for stdlib modules import side-effects
 # Defined globally to avoid adding one hook per test run (refleak)