Implement auto-import and evaluation policy overrides

krassowski · krassowski · commit d90a527a46f5 · 2025-05-30T23:40:02.000+01:00
diff --git a/IPython/core/completer.py b/IPython/core/completer.py
@@ -203,15 +203,11 @@
 from typing import (
     Iterable,
     Iterator,
-    List,
-    Tuple,
     Union,
     Any,
     Sequence,
-    Dict,
     Optional,
     TYPE_CHECKING,
-    Set,
     Sized,
     TypeVar,
     Literal,
@@ -236,10 +232,12 @@
     List as ListTrait,
     Unicode,
     Dict as DictTrait,
+    DottedObjectName,
     Union as UnionTrait,
     observe,
 )
 from traitlets.config.configurable import Configurable
+from traitlets.utils.importstring import import_item
 
 import __main__
 
@@ -987,15 +985,17 @@ class Completer(Configurable):
 
         - ``forbidden``: no evaluation of code is permitted,
         - ``minimal``: evaluation of literals and access to built-in namespace;
-          no item/attribute evaluationm no access to locals/globals,
+          no item/attribute evaluation, no access to locals/globals,
           no evaluation of any operations or comparisons.
         - ``limited``: access to all namespaces, evaluation of hard-coded methods
           (for example: :any:`dict.keys`, :any:`object.__getattr__`,
           :any:`object.__getitem__`) on allow-listed objects (for example:
           :any:`dict`, :any:`list`, :any:`tuple`, ``pandas.Series``),
         - ``unsafe``: evaluation of all methods and function calls but not of
           syntax with side-effects like `del x`,
-        - ``dangerous``: completely arbitrary evaluation.
+        - ``dangerous``: completely arbitrary evaluation; does not support auto-import.
+
+        To override specific elements of the policy, you can use ``policy_overrides`` trait.
         """,
     ).tag(config=True)
 
@@ -1031,6 +1031,35 @@ class Completer(Configurable):
         """,
     ).tag(config=True)
 
+    policy_overrides = DictTrait(
+        default_value={},
+        key_trait=Unicode(),
+        help="""Overrides for policy evaluation.
+
+        For example, to enable auto-import on completion specify:
+
+        .. code-block::
+
+            ipython --Completer.policy_overrides='{"allow_auto_import": True}' --Completer.use_jedi=False
+
+        """,
+    ).tag(config=True)
+
+    auto_import_method = DottedObjectName(
+        default_value="importlib.import_module",
+        allow_none=True,
+        help="""\
+        Provisional:
+              This is a provisional API in IPython 9.3, it may change without warnings.
+
+        A fully qualified path to an auto-import method for use by completer.
+        The function should take a single string and return `ModuleType` and
+        can raise `ImportError` exception if module is not found.
+
+        The default auto-import implementation does not populate the user namespace with the imported module.
+        """,
+    ).tag(config=True)
+
     def __init__(self, namespace=None, global_namespace=None, **kwargs):
         """Create a new completer for the command line.
 
@@ -1279,6 +1308,8 @@ def _evaluate_expr(self, expr):
                         globals=self.global_namespace,
                         locals=self.namespace,
                         evaluation=self.evaluation,
+                        auto_import=self._auto_import,
+                        policy_overrides=self.policy_overrides,
                     ),
                 )
                 done = True
@@ -1292,6 +1323,14 @@ def _evaluate_expr(self, expr):
                 expr = self._trim_expr(expr)
         return obj
 
+    @property
+    def _auto_import(self):
+        if self.auto_import_method is None:
+            return None
+        if not hasattr(self, "_auto_import_func"):
+            self._auto_import_func = import_item(self.auto_import_method)
+        return self._auto_import_func
+
 def get__all__entries(obj):
     """returns the strings in the __all__ attribute"""
     try:
@@ -2832,6 +2871,8 @@ def dict_key_matches(self, text: str) -> list[str]:
                 locals=self.namespace,
                 evaluation=self.evaluation,  # type: ignore
                 in_subscript=True,
+                auto_import=self._auto_import,
+                policy_overrides=self.policy_overrides,
             ),
         )
 
diff --git a/IPython/core/guarded_eval.py b/IPython/core/guarded_eval.py
@@ -1,19 +1,15 @@
+from copy import copy
 from inspect import isclass, signature, Signature
-from pygments.formatters.terminal256 import Terminal256Formatter
 from typing import (
     Annotated,
     AnyStr,
     Callable,
-    Dict,
     Literal,
     NamedTuple,
     NewType,
     Optional,
     Protocol,
-    Set,
     Sequence,
-    Tuple,
-    Type,
     TypeGuard,
     Union,
     get_args,
@@ -97,6 +93,7 @@ class EvaluationPolicy:
     allow_builtins_access: bool = False
     allow_all_operations: bool = False
     allow_any_calls: bool = False
+    allow_auto_import: bool = False
     allowed_calls: set[Callable] = field(default_factory=set)
 
     def can_get_item(self, value, item):
@@ -330,6 +327,10 @@ class EvaluationContext(NamedTuple):
     #: Whether the evaluation of code takes place inside of a subscript.
     #: Useful for evaluating ``:-1, 'col'`` in ``df[:-1, 'col']``.
     in_subscript: bool = False
+    #: Auto import method
+    auto_import: Callable[list[str], ModuleType] | None = None
+    #: Overrides for evaluation policy
+    policy_overrides: dict = {}
 
 
 class _IdentitySubscript:
@@ -463,6 +464,17 @@ def _find_dunder(node_op, dunders) -> Union[tuple[str, ...], None]:
     return dunder
 
 
+def get_policy(context: EvaluationContext) -> EvaluationPolicy:
+    policy = copy(EVALUATION_POLICIES[context.evaluation])
+
+    for key, value in context.policy_overrides.items():
+        if hasattr(policy, key):
+            setattr(policy, key, value)
+        else:
+            print(f"Incorrect policy override key: {key}")
+    return policy
+
+
 def eval_node(node: Union[ast.AST, None], context: EvaluationContext):
     """Evaluate AST node in provided context.
 
@@ -490,7 +502,8 @@ def eval_node(node: Union[ast.AST, None], context: EvaluationContext):
     The purpose of this function is to guard against unwanted side-effects;
     it does not give guarantees on protection from malicious code execution.
     """
-    policy = EVALUATION_POLICIES[context.evaluation]
+    policy = get_policy(context)
+
     if node is None:
         return None
     if isinstance(node, ast.Expression):
@@ -711,14 +724,16 @@ def _resolve_annotation(
 
 
 def _eval_node_name(node_id: str, context: EvaluationContext):
-    policy = EVALUATION_POLICIES[context.evaluation]
+    policy = get_policy(context)
     if policy.allow_locals_access and node_id in context.locals:
         return context.locals[node_id]
     if policy.allow_globals_access and node_id in context.globals:
         return context.globals[node_id]
     if policy.allow_builtins_access and hasattr(builtins, node_id):
         # note: do not use __builtins__, it is implementation detail of cPython
         return getattr(builtins, node_id)
+    if policy.allow_auto_import and context.auto_import:
+        return context.auto_import(node_id)
     if not policy.allow_globals_access and not policy.allow_locals_access:
         raise GuardRejection(
             f"Namespace access not allowed in {context.evaluation} mode"
@@ -728,7 +743,7 @@ def _eval_node_name(node_id: str, context: EvaluationContext):
 
 
 def _eval_or_create_duck(duck_type, node: ast.Call, context: EvaluationContext):
-    policy = EVALUATION_POLICIES[context.evaluation]
+    policy = get_policy(context)
     # if allow-listed builtin is on type annotation, instantiate it
     if policy.can_call(duck_type) and not node.keywords:
         args = [eval_node(arg, context) for arg in node.args]
diff --git a/IPython/terminal/interactiveshell.py b/IPython/terminal/interactiveshell.py
@@ -475,7 +475,7 @@ def _llm_prefix_from_history_changed(self, change):
         allow_none=True,
         help="""\
         Provisional:
-            This is a provisinal API in IPython 8.32, before stabilisation
+            This is a provisional API in IPython 8.32, before stabilisation
             in 9.0, it may change without warnings.
 
         class to use for the `NavigableAutoSuggestFromHistory` to request
diff --git a/tests/test_completer.py b/tests/test_completer.py
@@ -141,14 +141,17 @@ def greedy_completion():
 
 
 @contextmanager
-def evaluation_policy(evaluation: str):
+def evaluation_policy(evaluation: str, **overrides):
     ip = get_ipython()
     evaluation_original = ip.Completer.evaluation
+    overrides_original = ip.Completer.policy_overrides
     try:
         ip.Completer.evaluation = evaluation
+        ip.Completer.policy_overrides = overrides
         yield
     finally:
         ip.Completer.evaluation = evaluation_original
+        ip.Completer.policy_overrides = overrides_original
 
 
 @contextmanager
@@ -1380,6 +1383,26 @@ def assert_completion(**kwargs):
             assert_completion(line_buffer="get()['ab")
             assert_completion(line_buffer="get()['abc")
 
+    def test_completion_autoimport(self):
+        ip = get_ipython()
+        complete = ip.Completer.complete
+        with (
+            evaluation_policy("limited", allow_auto_import=True),
+            jedi_status(False),
+        ):
+            _, matches = complete(line_buffer="math.")
+            self.assertIn(".pi", matches)
+
+    def test_completion_no_autoimport(self):
+        ip = get_ipython()
+        complete = ip.Completer.complete
+        with (
+            evaluation_policy("limited", allow_auto_import=False),
+            jedi_status(False),
+        ):
+            _, matches = complete(line_buffer="math.")
+            self.assertNotIn(".pi", matches)
+
     def test_dict_key_completion_bytes(self):
         """Test handling of bytes in dict key completion"""
         ip = get_ipython()
diff --git a/tests/test_guarded_eval.py b/tests/test_guarded_eval.py
@@ -1,5 +1,6 @@
 import sys
 from contextlib import contextmanager
+from importlib import import_module
 from typing import (
     Annotated,
     AnyStr,
@@ -780,3 +781,19 @@ def test_module_access():
     context = minimal(numpy=numpy)
     with pytest.raises(GuardRejection):
         guarded_eval("np.linalg.norm", context)
+
+
+def test_autoimport_module():
+    context = EvaluationContext(
+        locals={}, globals={}, evaluation="limited", auto_import=import_module
+    )
+    pi = guarded_eval("math.pi", context)
+    assert round(pi, 2) == 3.14
+
+
+def test_autoimport_deep_module():
+    context = EvaluationContext(
+        locals={}, globals={}, evaluation="limited", auto_import=import_module
+    )
+    ElementTree = guarded_eval("xml.etree.ElementTree", context)
+    assert hasattr(ElementTree, "ElementTree")
