Support stringized return type annotations

IPython/core/guarded_eval.py

@@ -1,793 +1,803 b''
 from inspect import isclass, signature, Signature
 from typing import (
-    Any,
     Callable,
     Dict,
     Set,
     Sequence,
     Tuple,
     NamedTuple,
     Type,
     Literal,
     Union,
     TYPE_CHECKING,
 )
 import ast
 import builtins
 import collections
 import operator
 import sys
 from functools import cached_property
 from dataclasses import dataclass, field
 from types import MethodDescriptorType, ModuleType
 
 from IPython.utils.docs import GENERATING_DOCUMENTATION
 from IPython.utils.decorators import undoc
 
 
 if TYPE_CHECKING or GENERATING_DOCUMENTATION:
     from typing_extensions import Protocol
 else:
     # do not require on runtime
     Protocol = object  # requires Python >=3.8
 
 
 @undoc
 class HasGetItem(Protocol):
     def __getitem__(self, key) -> None:
         ...
 
 
 @undoc
 class InstancesHaveGetItem(Protocol):
     def __call__(self, *args, **kwargs) -> HasGetItem:
         ...
 
 
 @undoc
 class HasGetAttr(Protocol):
     def __getattr__(self, key) -> None:
         ...
 
 
 @undoc
 class DoesNotHaveGetAttr(Protocol):
     pass
 
 
 # By default `__getattr__` is not explicitly implemented on most objects
 MayHaveGetattr = Union[HasGetAttr, DoesNotHaveGetAttr]
 
 
 def _unbind_method(func: Callable) -> Union[Callable, None]:
     """Get unbound method for given bound method.
 
     Returns None if cannot get unbound method, or method is already unbound.
     """
     owner = getattr(func, "__self__", None)
     owner_class = type(owner)
     name = getattr(func, "__name__", None)
     instance_dict_overrides = getattr(owner, "__dict__", None)
     if (
         owner is not None
         and name
         and (
             not instance_dict_overrides
             or (instance_dict_overrides and name not in instance_dict_overrides)
         )
     ):
         return getattr(owner_class, name)
     return None
 
 
 @undoc
 @dataclass
 class EvaluationPolicy:
     """Definition of evaluation policy."""
 
     allow_locals_access: bool = False
     allow_globals_access: bool = False
     allow_item_access: bool = False
     allow_attr_access: bool = False
     allow_builtins_access: bool = False
     allow_all_operations: bool = False
     allow_any_calls: bool = False
     allowed_calls: Set[Callable] = field(default_factory=set)
 
     def can_get_item(self, value, item):
         return self.allow_item_access
 
     def can_get_attr(self, value, attr):
         return self.allow_attr_access
 
     def can_operate(self, dunders: Tuple[str, ...], a, b=None):
         if self.allow_all_operations:
             return True
 
     def can_call(self, func):
         if self.allow_any_calls:
             return True
 
         if func in self.allowed_calls:
             return True
 
         owner_method = _unbind_method(func)
 
         if owner_method and owner_method in self.allowed_calls:
             return True
 
 
 def _get_external(module_name: str, access_path: Sequence[str]):
     """Get value from external module given a dotted access path.
 
     Raises:
     * `KeyError` if module is removed not found, and
     * `AttributeError` if acess path does not match an exported object
     """
     member_type = sys.modules[module_name]
     for attr in access_path:
         member_type = getattr(member_type, attr)
     return member_type
 
 
 def _has_original_dunder_external(
     value,
     module_name: str,
     access_path: Sequence[str],
     method_name: str,
 ):
     if module_name not in sys.modules:
         # LBYLB as it is faster
         return False
     try:
         member_type = _get_external(module_name, access_path)
         value_type = type(value)
         if type(value) == member_type:
             return True
         if method_name == "__getattribute__":
             # we have to short-circuit here due to an unresolved issue in
             # `isinstance` implementation: https://bugs.python.org/issue32683
             return False
         if isinstance(value, member_type):
             method = getattr(value_type, method_name, None)
             member_method = getattr(member_type, method_name, None)
             if member_method == method:
                 return True
     except (AttributeError, KeyError):
         return False
 
 
 def _has_original_dunder(
     value, allowed_types, allowed_methods, allowed_external, method_name
 ):
     # note: Python ignores `__getattr__`/`__getitem__` on instances,
     # we only need to check at class level
     value_type = type(value)
 
     # strict type check passes → no need to check method
     if value_type in allowed_types:
         return True
 
     method = getattr(value_type, method_name, None)
 
     if method is None:
         return None
 
     if method in allowed_methods:
         return True
 
     for module_name, *access_path in allowed_external:
         if _has_original_dunder_external(value, module_name, access_path, method_name):
             return True
 
     return False
 
 
 @undoc
 @dataclass
 class SelectivePolicy(EvaluationPolicy):
     allowed_getitem: Set[InstancesHaveGetItem] = field(default_factory=set)
     allowed_getitem_external: Set[Tuple[str, ...]] = field(default_factory=set)
 
     allowed_getattr: Set[MayHaveGetattr] = field(default_factory=set)
     allowed_getattr_external: Set[Tuple[str, ...]] = field(default_factory=set)
 
     allowed_operations: Set = field(default_factory=set)
     allowed_operations_external: Set[Tuple[str, ...]] = field(default_factory=set)
 
     _operation_methods_cache: Dict[str, Set[Callable]] = field(
         default_factory=dict, init=False
     )
 
     def can_get_attr(self, value, attr):
         has_original_attribute = _has_original_dunder(
             value,
             allowed_types=self.allowed_getattr,
             allowed_methods=self._getattribute_methods,
             allowed_external=self.allowed_getattr_external,
             method_name="__getattribute__",
         )
         has_original_attr = _has_original_dunder(
             value,
             allowed_types=self.allowed_getattr,
             allowed_methods=self._getattr_methods,
             allowed_external=self.allowed_getattr_external,
             method_name="__getattr__",
         )
 
         accept = False
 
         # Many objects do not have `__getattr__`, this is fine.
         if has_original_attr is None and has_original_attribute:
             accept = True
         else:
             # Accept objects without modifications to `__getattr__` and `__getattribute__`
             accept = has_original_attr and has_original_attribute
 
         if accept:
             # We still need to check for overriden properties.
 
             value_class = type(value)
             if not hasattr(value_class, attr):
                 return True
 
             class_attr_val = getattr(value_class, attr)
             is_property = isinstance(class_attr_val, property)
 
             if not is_property:
                 return True
 
             # Properties in allowed types are ok (although we do not include any
             # properties in our default allow list currently).
             if type(value) in self.allowed_getattr:
                 return True  # pragma: no cover
 
             # Properties in subclasses of allowed types may be ok if not changed
             for module_name, *access_path in self.allowed_getattr_external:
                 try:
                     external_class = _get_external(module_name, access_path)
                     external_class_attr_val = getattr(external_class, attr)
                 except (KeyError, AttributeError):
                     return False  # pragma: no cover
                 return class_attr_val == external_class_attr_val
 
         return False
 
     def can_get_item(self, value, item):
         """Allow accessing `__getiitem__` of allow-listed instances unless it was not modified."""
         return _has_original_dunder(
             value,
             allowed_types=self.allowed_getitem,
             allowed_methods=self._getitem_methods,
             allowed_external=self.allowed_getitem_external,
             method_name="__getitem__",
         )
 
     def can_operate(self, dunders: Tuple[str, ...], a, b=None):
         objects = [a]
         if b is not None:
             objects.append(b)
         return all(
             [
                 _has_original_dunder(
                     obj,
                     allowed_types=self.allowed_operations,
                     allowed_methods=self._operator_dunder_methods(dunder),
                     allowed_external=self.allowed_operations_external,
                     method_name=dunder,
                 )
                 for dunder in dunders
                 for obj in objects
             ]
         )
 
     def _operator_dunder_methods(self, dunder: str) -> Set[Callable]:
         if dunder not in self._operation_methods_cache:
             self._operation_methods_cache[dunder] = self._safe_get_methods(
                 self.allowed_operations, dunder
             )
         return self._operation_methods_cache[dunder]
 
     @cached_property
     def _getitem_methods(self) -> Set[Callable]:
         return self._safe_get_methods(self.allowed_getitem, "__getitem__")
 
     @cached_property
     def _getattr_methods(self) -> Set[Callable]:
         return self._safe_get_methods(self.allowed_getattr, "__getattr__")
 
     @cached_property
     def _getattribute_methods(self) -> Set[Callable]:
         return self._safe_get_methods(self.allowed_getattr, "__getattribute__")
 
     def _safe_get_methods(self, classes, name) -> Set[Callable]:
         return {
             method
             for class_ in classes
             for method in [getattr(class_, name, None)]
             if method
         }
 
 
 class _DummyNamedTuple(NamedTuple):
     """Used internally to retrieve methods of named tuple instance."""
 
 
 class EvaluationContext(NamedTuple):
     #: Local namespace
     locals: dict
     #: Global namespace
     globals: dict
     #: Evaluation policy identifier
     evaluation: Literal[
         "forbidden", "minimal", "limited", "unsafe", "dangerous"
     ] = "forbidden"
     #: Whether the evalution of code takes place inside of a subscript.
     #: Useful for evaluating ``:-1, 'col'`` in ``df[:-1, 'col']``.
     in_subscript: bool = False
 
 
 class _IdentitySubscript:
     """Returns the key itself when item is requested via subscript."""
 
     def __getitem__(self, key):
         return key
 
 
 IDENTITY_SUBSCRIPT = _IdentitySubscript()
 SUBSCRIPT_MARKER = "__SUBSCRIPT_SENTINEL__"
 UNKNOWN_SIGNATURE = Signature()
 NOT_EVALUATED = object()
 
 
 class GuardRejection(Exception):
     """Exception raised when guard rejects evaluation attempt."""
 
     pass
 
 
 def guarded_eval(code: str, context: EvaluationContext):
     """Evaluate provided code in the evaluation context.
 
     If evaluation policy given by context is set to ``forbidden``
     no evaluation will be performed; if it is set to ``dangerous``
     standard :func:`eval` will be used; finally, for any other,
     policy :func:`eval_node` will be called on parsed AST.
     """
     locals_ = context.locals
 
     if context.evaluation == "forbidden":
         raise GuardRejection("Forbidden mode")
 
     # note: not using `ast.literal_eval` as it does not implement
     # getitem at all, for example it fails on simple `[0][1]`
 
     if context.in_subscript:
         # syntatic sugar for ellipsis (:) is only available in susbcripts
         # so we need to trick the ast parser into thinking that we have
         # a subscript, but we need to be able to later recognise that we did
         # it so we can ignore the actual __getitem__ operation
         if not code:
             return tuple()
         locals_ = locals_.copy()
         locals_[SUBSCRIPT_MARKER] = IDENTITY_SUBSCRIPT
         code = SUBSCRIPT_MARKER + "[" + code + "]"
         context = EvaluationContext(**{**context._asdict(), **{"locals": locals_}})
 
     if context.evaluation == "dangerous":
         return eval(code, context.globals, context.locals)
 
     expression = ast.parse(code, mode="eval")
 
     return eval_node(expression, context)
 
 
 BINARY_OP_DUNDERS: Dict[Type[ast.operator], Tuple[str]] = {
     ast.Add: ("__add__",),
     ast.Sub: ("__sub__",),
     ast.Mult: ("__mul__",),
     ast.Div: ("__truediv__",),
     ast.FloorDiv: ("__floordiv__",),
     ast.Mod: ("__mod__",),
     ast.Pow: ("__pow__",),
     ast.LShift: ("__lshift__",),
     ast.RShift: ("__rshift__",),
     ast.BitOr: ("__or__",),
     ast.BitXor: ("__xor__",),
     ast.BitAnd: ("__and__",),
     ast.MatMult: ("__matmul__",),
 }
 
 COMP_OP_DUNDERS: Dict[Type[ast.cmpop], Tuple[str, ...]] = {
     ast.Eq: ("__eq__",),
     ast.NotEq: ("__ne__", "__eq__"),
     ast.Lt: ("__lt__", "__gt__"),
     ast.LtE: ("__le__", "__ge__"),
     ast.Gt: ("__gt__", "__lt__"),
     ast.GtE: ("__ge__", "__le__"),
     ast.In: ("__contains__",),
     # Note: ast.Is, ast.IsNot, ast.NotIn are handled specially
 }
 
 UNARY_OP_DUNDERS: Dict[Type[ast.unaryop], Tuple[str, ...]] = {
     ast.USub: ("__neg__",),
     ast.UAdd: ("__pos__",),
     # we have to check both __inv__ and __invert__!
     ast.Invert: ("__invert__", "__inv__"),
     ast.Not: ("__not__",),
 }
 
 
 class Duck:
     """A dummy class used to create objects of other classes without calling their ``__init__``"""
 
 
 def _find_dunder(node_op, dunders) -> Union[Tuple[str, ...], None]:
     dunder = None
     for op, candidate_dunder in dunders.items():
         if isinstance(node_op, op):
             dunder = candidate_dunder
     return dunder
 
 
 def eval_node(node: Union[ast.AST, None], context: EvaluationContext):
     """Evaluate AST node in provided context.
 
     Applies evaluation restrictions defined in the context. Currently does not support evaluation of functions with keyword arguments.
 
     Does not evaluate actions that always have side effects:
 
     - class definitions (``class sth: ...``)
     - function definitions (``def sth: ...``)
     - variable assignments (``x = 1``)
     - augmented assignments (``x += 1``)
     - deletions (``del x``)
 
     Does not evaluate operations which do not return values:
 
     - assertions (``assert x``)
     - pass (``pass``)
     - imports (``import x``)
     - control flow:
 
         - conditionals (``if x:``) except for ternary IfExp (``a if x else b``)
         - loops (``for`` and ``while``)
         - exception handling
 
     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]
     if node is None:
         return None
     if isinstance(node, ast.Expression):
         return eval_node(node.body, context)
     if isinstance(node, ast.BinOp):
         left = eval_node(node.left, context)
         right = eval_node(node.right, context)
         dunders = _find_dunder(node.op, BINARY_OP_DUNDERS)
         if dunders:
             if policy.can_operate(dunders, left, right):
                 return getattr(left, dunders[0])(right)
             else:
                 raise GuardRejection(
                     f"Operation (`{dunders}`) for",
                     type(left),
                     f"not allowed in {context.evaluation} mode",
                 )
     if isinstance(node, ast.Compare):
         left = eval_node(node.left, context)
         all_true = True
         negate = False
         for op, right in zip(node.ops, node.comparators):
             right = eval_node(right, context)
             dunder = None
             dunders = _find_dunder(op, COMP_OP_DUNDERS)
             if not dunders:
                 if isinstance(op, ast.NotIn):
                     dunders = COMP_OP_DUNDERS[ast.In]
                     negate = True
                 if isinstance(op, ast.Is):
                     dunder = "is_"
                 if isinstance(op, ast.IsNot):
                     dunder = "is_"
                     negate = True
             if not dunder and dunders:
                 dunder = dunders[0]
             if dunder:
                 a, b = (right, left) if dunder == "__contains__" else (left, right)
                 if dunder == "is_" or dunders and policy.can_operate(dunders, a, b):
                     result = getattr(operator, dunder)(a, b)
                     if negate:
                         result = not result
                     if not result:
                         all_true = False
                     left = right
                 else:
                     raise GuardRejection(
                         f"Comparison (`{dunder}`) for",
                         type(left),
                         f"not allowed in {context.evaluation} mode",
                     )
             else:
                 raise ValueError(
                     f"Comparison `{dunder}` not supported"
                 )  # pragma: no cover
         return all_true
     if isinstance(node, ast.Constant):
         return node.value
     if isinstance(node, ast.Tuple):
         return tuple(eval_node(e, context) for e in node.elts)
     if isinstance(node, ast.List):
         return [eval_node(e, context) for e in node.elts]
     if isinstance(node, ast.Set):
         return {eval_node(e, context) for e in node.elts}
     if isinstance(node, ast.Dict):
         return dict(
             zip(
                 [eval_node(k, context) for k in node.keys],
                 [eval_node(v, context) for v in node.values],
             )
         )
     if isinstance(node, ast.Slice):
         return slice(
             eval_node(node.lower, context),
             eval_node(node.upper, context),
             eval_node(node.step, context),
         )
     if isinstance(node, ast.UnaryOp):
         value = eval_node(node.operand, context)
         dunders = _find_dunder(node.op, UNARY_OP_DUNDERS)
         if dunders:
             if policy.can_operate(dunders, value):
                 return getattr(value, dunders[0])()
             else:
                 raise GuardRejection(
                     f"Operation (`{dunders}`) for",
                     type(value),
                     f"not allowed in {context.evaluation} mode",
                 )
     if isinstance(node, ast.Subscript):
         value = eval_node(node.value, context)
         slice_ = eval_node(node.slice, context)
         if policy.can_get_item(value, slice_):
             return value[slice_]
         raise GuardRejection(
             "Subscript access (`__getitem__`) for",
             type(value),  # not joined to avoid calling `repr`
             f" not allowed in {context.evaluation} mode",
         )
     if isinstance(node, ast.Name):
-        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 not policy.allow_globals_access and not policy.allow_locals_access:
-            raise GuardRejection(
-                f"Namespace access not allowed in {context.evaluation} mode"
-            )
-        else:
-            raise NameError(f"{node.id} not found in locals, globals, nor builtins")
+        return _eval_node_name(node.id, policy, context)
     if isinstance(node, ast.Attribute):
         value = eval_node(node.value, context)
         if policy.can_get_attr(value, node.attr):
             return getattr(value, node.attr)
         raise GuardRejection(
             "Attribute access (`__getattr__`) for",
             type(value),  # not joined to avoid calling `repr`
             f"not allowed in {context.evaluation} mode",
         )
     if isinstance(node, ast.IfExp):
         test = eval_node(node.test, context)
         if test:
             return eval_node(node.body, context)
         else:
             return eval_node(node.orelse, context)
     if isinstance(node, ast.Call):
         func = eval_node(node.func, context)
         if policy.can_call(func) and not node.keywords:
             args = [eval_node(arg, context) for arg in node.args]
             return func(*args)
         if isclass(func):
             # this code path gets entered when calling class e.g. `MyClass()`
             # or `my_instance.__class__()` - in both cases `func` is `MyClass`.
             # Should return `MyClass` if `__new__` is not overridden,
             # otherwise whatever `__new__` return type is.
             overridden_return_type = _eval_return_type(
                 func.__new__, policy, node, context
             )
             if overridden_return_type is not NOT_EVALUATED:
                 return overridden_return_type
             return _create_duck_for_type(func)
         else:
             return_type = _eval_return_type(func, policy, node, context)
             if return_type is not NOT_EVALUATED:
                 return return_type
         raise GuardRejection(
             "Call for",
             func,  # not joined to avoid calling `repr`
             f"not allowed in {context.evaluation} mode",
         )
     raise ValueError("Unhandled node", ast.dump(node))
 
 
-def _eval_return_type(func, policy, node, context):
+def _eval_return_type(
+    func: Callable, policy: EvaluationPolicy, node: ast.Call, context: EvaluationContext
+):
     """Evaluate return type of a given callable function.
 
     Returns the built-in type, a duck or NOT_EVALUATED sentinel.
     """
     try:
         sig = signature(func)
     except ValueError:
         sig = UNKNOWN_SIGNATURE
     # if annotation was not stringized, or it was stringized
     # but resolved by signature call we know the return type
     not_empty = sig.return_annotation is not Signature.empty
-    not_stringized = not isinstance(sig.return_annotation, str)
-    if not_empty and not_stringized:
+    stringized = isinstance(sig.return_annotation, str)
+    if not_empty:
+        return_type = (
+            _eval_node_name(sig.return_annotation, policy, context)
+            if stringized
+            else sig.return_annotation
+        )
         # if allow-listed builtin is on type annotation, instantiate it
-        if policy.can_call(sig.return_annotation) and not node.keywords:
+        if policy.can_call(return_type) and not node.keywords:
             args = [eval_node(arg, context) for arg in node.args]
             # if custom class is in type annotation, mock it;
-            return sig.return_annotation(*args)
-        return _create_duck_for_type(sig.return_annotation)
+            return return_type(*args)
+        return _create_duck_for_type(return_type)
     return NOT_EVALUATED
 
 
+def _eval_node_name(node_id: str, policy: EvaluationPolicy, context: EvaluationContext):
+    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 not policy.allow_globals_access and not policy.allow_locals_access:
+        raise GuardRejection(
+            f"Namespace access not allowed in {context.evaluation} mode"
+        )
+    else:
+        raise NameError(f"{node_id} not found in locals, globals, nor builtins")
+
+
 def _create_duck_for_type(duck_type):
     """Create an imitation of an object of a given type (a duck).
 
     Returns the duck or NOT_EVALUATED sentinel if duck could not be created.
     """
     duck = Duck()
     try:
         # this only works for heap types, not builtins
         duck.__class__ = duck_type
         return duck
     except TypeError:
         pass
     return NOT_EVALUATED
 
 
 SUPPORTED_EXTERNAL_GETITEM = {
     ("pandas", "core", "indexing", "_iLocIndexer"),
     ("pandas", "core", "indexing", "_LocIndexer"),
     ("pandas", "DataFrame"),
     ("pandas", "Series"),
     ("numpy", "ndarray"),
     ("numpy", "void"),
 }
 
 
 BUILTIN_GETITEM: Set[InstancesHaveGetItem] = {
     dict,
     str,  # type: ignore[arg-type]
     bytes,  # type: ignore[arg-type]
     list,
     tuple,
     collections.defaultdict,
     collections.deque,
     collections.OrderedDict,
     collections.ChainMap,
     collections.UserDict,
     collections.UserList,
     collections.UserString,  # type: ignore[arg-type]
     _DummyNamedTuple,
     _IdentitySubscript,
 }
 
 
 def _list_methods(cls, source=None):
     """For use on immutable objects or with methods returning a copy"""
     return [getattr(cls, k) for k in (source if source else dir(cls))]
 
 
 dict_non_mutating_methods = ("copy", "keys", "values", "items")
 list_non_mutating_methods = ("copy", "index", "count")
 set_non_mutating_methods = set(dir(set)) & set(dir(frozenset))
 
 
 dict_keys: Type[collections.abc.KeysView] = type({}.keys())
 
 NUMERICS = {int, float, complex}
 
 ALLOWED_CALLS = {
     bytes,
     *_list_methods(bytes),
     dict,
     *_list_methods(dict, dict_non_mutating_methods),
     dict_keys.isdisjoint,
     list,
     *_list_methods(list, list_non_mutating_methods),
     set,
     *_list_methods(set, set_non_mutating_methods),
     frozenset,
     *_list_methods(frozenset),
     range,
     str,
     *_list_methods(str),
     tuple,
     *_list_methods(tuple),
     *NUMERICS,
     *[method for numeric_cls in NUMERICS for method in _list_methods(numeric_cls)],
     collections.deque,
     *_list_methods(collections.deque, list_non_mutating_methods),
     collections.defaultdict,
     *_list_methods(collections.defaultdict, dict_non_mutating_methods),
     collections.OrderedDict,
     *_list_methods(collections.OrderedDict, dict_non_mutating_methods),
     collections.UserDict,
     *_list_methods(collections.UserDict, dict_non_mutating_methods),
     collections.UserList,
     *_list_methods(collections.UserList, list_non_mutating_methods),
     collections.UserString,
     *_list_methods(collections.UserString, dir(str)),
     collections.Counter,
     *_list_methods(collections.Counter, dict_non_mutating_methods),
     collections.Counter.elements,
     collections.Counter.most_common,
 }
 
 BUILTIN_GETATTR: Set[MayHaveGetattr] = {
     *BUILTIN_GETITEM,
     set,
     frozenset,
     object,
     type,  # `type` handles a lot of generic cases, e.g. numbers as in `int.real`.
     *NUMERICS,
     dict_keys,
     MethodDescriptorType,
     ModuleType,
 }
 
 
 BUILTIN_OPERATIONS = {*BUILTIN_GETATTR}
 
 EVALUATION_POLICIES = {
     "minimal": EvaluationPolicy(
         allow_builtins_access=True,
         allow_locals_access=False,
         allow_globals_access=False,
         allow_item_access=False,
         allow_attr_access=False,
         allowed_calls=set(),
         allow_any_calls=False,
         allow_all_operations=False,
     ),
     "limited": SelectivePolicy(
         allowed_getitem=BUILTIN_GETITEM,
         allowed_getitem_external=SUPPORTED_EXTERNAL_GETITEM,
         allowed_getattr=BUILTIN_GETATTR,
         allowed_getattr_external={
             # pandas Series/Frame implements custom `__getattr__`
             ("pandas", "DataFrame"),
             ("pandas", "Series"),
         },
         allowed_operations=BUILTIN_OPERATIONS,
         allow_builtins_access=True,
         allow_locals_access=True,
         allow_globals_access=True,
         allowed_calls=ALLOWED_CALLS,
     ),
     "unsafe": EvaluationPolicy(
         allow_builtins_access=True,
         allow_locals_access=True,
         allow_globals_access=True,
         allow_attr_access=True,
         allow_item_access=True,
         allow_any_calls=True,
         allow_all_operations=True,
     ),
 }
 
 
 __all__ = [
     "guarded_eval",
     "eval_node",
     "GuardRejection",
     "EvaluationContext",
     "_unbind_method",
 ]

IPython/core/tests/test_guarded_eval.py

@@ -1,631 +1,641 b''
 from contextlib import contextmanager
 from typing import NamedTuple
 from functools import partial
 from IPython.core.guarded_eval import (
     EvaluationContext,
     GuardRejection,
     guarded_eval,
     _unbind_method,
 )
 from IPython.testing import decorators as dec
 import pytest
 
 
 def create_context(evaluation: str, **kwargs):
     return EvaluationContext(locals=kwargs, globals={}, evaluation=evaluation)
 
 
 forbidden = partial(create_context, "forbidden")
 minimal = partial(create_context, "minimal")
 limited = partial(create_context, "limited")
 unsafe = partial(create_context, "unsafe")
 dangerous = partial(create_context, "dangerous")
 
 LIMITED_OR_HIGHER = [limited, unsafe, dangerous]
 MINIMAL_OR_HIGHER = [minimal, *LIMITED_OR_HIGHER]
 
 
 @contextmanager
 def module_not_installed(module: str):
     import sys
 
     try:
         to_restore = sys.modules[module]
         del sys.modules[module]
     except KeyError:
         to_restore = None
     try:
         yield
     finally:
         sys.modules[module] = to_restore
 
 
 def test_external_not_installed():
     """
     Because attribute check requires checking if object is not of allowed
     external type, this tests logic for absence of external module.
     """
 
     class Custom:
         def __init__(self):
             self.test = 1
 
         def __getattr__(self, key):
             return key
 
     with module_not_installed("pandas"):
         context = limited(x=Custom())
         with pytest.raises(GuardRejection):
             guarded_eval("x.test", context)
 
 
 @dec.skip_without("pandas")
 def test_external_changed_api(monkeypatch):
     """Check that the execution rejects if external API changed paths"""
     import pandas as pd
 
     series = pd.Series([1], index=["a"])
 
     with monkeypatch.context() as m:
         m.delattr(pd, "Series")
         context = limited(data=series)
         with pytest.raises(GuardRejection):
             guarded_eval("data.iloc[0]", context)
 
 
 @dec.skip_without("pandas")
 def test_pandas_series_iloc():
     import pandas as pd
 
     series = pd.Series([1], index=["a"])
     context = limited(data=series)
     assert guarded_eval("data.iloc[0]", context) == 1
 
 
 def test_rejects_custom_properties():
     class BadProperty:
         @property
         def iloc(self):
             return [None]
 
     series = BadProperty()
     context = limited(data=series)
 
     with pytest.raises(GuardRejection):
         guarded_eval("data.iloc[0]", context)
 
 
 @dec.skip_without("pandas")
 def test_accepts_non_overriden_properties():
     import pandas as pd
 
     class GoodProperty(pd.Series):
         pass
 
     series = GoodProperty([1], index=["a"])
     context = limited(data=series)
 
     assert guarded_eval("data.iloc[0]", context) == 1
 
 
 @dec.skip_without("pandas")
 def test_pandas_series():
     import pandas as pd
 
     context = limited(data=pd.Series([1], index=["a"]))
     assert guarded_eval('data["a"]', context) == 1
     with pytest.raises(KeyError):
         guarded_eval('data["c"]', context)
 
 
 @dec.skip_without("pandas")
 def test_pandas_bad_series():
     import pandas as pd
 
     class BadItemSeries(pd.Series):
         def __getitem__(self, key):
             return "CUSTOM_ITEM"
 
     class BadAttrSeries(pd.Series):
         def __getattr__(self, key):
             return "CUSTOM_ATTR"
 
     bad_series = BadItemSeries([1], index=["a"])
     context = limited(data=bad_series)
 
     with pytest.raises(GuardRejection):
         guarded_eval('data["a"]', context)
     with pytest.raises(GuardRejection):
         guarded_eval('data["c"]', context)
 
     # note: here result is a bit unexpected because
     # pandas `__getattr__` calls `__getitem__`;
     # FIXME - special case to handle it?
     assert guarded_eval("data.a", context) == "CUSTOM_ITEM"
 
     context = unsafe(data=bad_series)
     assert guarded_eval('data["a"]', context) == "CUSTOM_ITEM"
 
     bad_attr_series = BadAttrSeries([1], index=["a"])
     context = limited(data=bad_attr_series)
     assert guarded_eval('data["a"]', context) == 1
     with pytest.raises(GuardRejection):
         guarded_eval("data.a", context)
 
 
 @dec.skip_without("pandas")
 def test_pandas_dataframe_loc():
     import pandas as pd
     from pandas.testing import assert_series_equal
 
     data = pd.DataFrame([{"a": 1}])
     context = limited(data=data)
     assert_series_equal(guarded_eval('data.loc[:, "a"]', context), data["a"])
 
 
 def test_named_tuple():
     class GoodNamedTuple(NamedTuple):
         a: str
         pass
 
     class BadNamedTuple(NamedTuple):
         a: str
 
         def __getitem__(self, key):
             return None
 
     good = GoodNamedTuple(a="x")
     bad = BadNamedTuple(a="x")
 
     context = limited(data=good)
     assert guarded_eval("data[0]", context) == "x"
 
     context = limited(data=bad)
     with pytest.raises(GuardRejection):
         guarded_eval("data[0]", context)
 
 
 def test_dict():
     context = limited(data={"a": 1, "b": {"x": 2}, ("x", "y"): 3})
     assert guarded_eval('data["a"]', context) == 1
     assert guarded_eval('data["b"]', context) == {"x": 2}
     assert guarded_eval('data["b"]["x"]', context) == 2
     assert guarded_eval('data["x", "y"]', context) == 3
 
     assert guarded_eval("data.keys", context)
 
 
 def test_set():
     context = limited(data={"a", "b"})
     assert guarded_eval("data.difference", context)
 
 
 def test_list():
     context = limited(data=[1, 2, 3])
     assert guarded_eval("data[1]", context) == 2
     assert guarded_eval("data.copy", context)
 
 
 def test_dict_literal():
     context = limited()
     assert guarded_eval("{}", context) == {}
     assert guarded_eval('{"a": 1}', context) == {"a": 1}
 
 
 def test_list_literal():
     context = limited()
     assert guarded_eval("[]", context) == []
     assert guarded_eval('[1, "a"]', context) == [1, "a"]
 
 
 def test_set_literal():
     context = limited()
     assert guarded_eval("set()", context) == set()
     assert guarded_eval('{"a"}', context) == {"a"}
 
 
 def test_evaluates_if_expression():
     context = limited()
     assert guarded_eval("2 if True else 3", context) == 2
     assert guarded_eval("4 if False else 5", context) == 5
 
 
 def test_object():
     obj = object()
     context = limited(obj=obj)
     assert guarded_eval("obj.__dir__", context) == obj.__dir__
 
 
 @pytest.mark.parametrize(
     "code,expected",
     [
         ["int.numerator", int.numerator],
         ["float.is_integer", float.is_integer],
         ["complex.real", complex.real],
     ],
 )
 def test_number_attributes(code, expected):
     assert guarded_eval(code, limited()) == expected
 
 
 def test_method_descriptor():
     context = limited()
     assert guarded_eval("list.copy.__name__", context) == "copy"
 
 
 class HeapType:
     pass
 
 
 class CallCreatesHeapType:
     def __call__(self) -> HeapType:
         return HeapType()
 
 
 class CallCreatesBuiltin:
     def __call__(self) -> frozenset:
         return frozenset()
 
 
 class HasStaticMethod:
     @staticmethod
     def static_method() -> HeapType:
         return HeapType()
 
 
 class InitReturnsFrozenset:
     def __new__(self) -> frozenset:  # type:ignore[misc]
         return frozenset()
 
 
+class StringAnnotation:
+    def heap(self) -> "HeapType":
+        return HeapType()
+
+    def copy(self) -> "StringAnnotation":
+        return StringAnnotation()
+
+
 @pytest.mark.parametrize(
     "data,good,expected,equality",
     [
         [[1, 2, 3], "data.index(2)", 1, True],
         [{"a": 1}, "data.keys().isdisjoint({})", True, True],
+        [StringAnnotation(), "data.heap()", HeapType, False],
+        [StringAnnotation(), "data.copy()", StringAnnotation, False],
         # test cases for `__call__`
         [CallCreatesHeapType(), "data()", HeapType, False],
         [CallCreatesBuiltin(), "data()", frozenset, False],
         # Test cases for `__init__`
         [HeapType, "data()", HeapType, False],
         [InitReturnsFrozenset, "data()", frozenset, False],
         [HeapType(), "data.__class__()", HeapType, False],
-        # test cases for static and class methods
+        # test cases for static methods
         [HasStaticMethod, "data.static_method()", HeapType, False],
     ],
 )
 def test_evaluates_calls(data, good, expected, equality):
-    context = limited(data=data)
+    context = limited(data=data, HeapType=HeapType, StringAnnotation=StringAnnotation)
     value = guarded_eval(good, context)
     if equality:
         assert value == expected
     else:
         assert isinstance(value, expected)
 
 
 @pytest.mark.parametrize(
     "data,bad",
     [
         [[1, 2, 3], "data.append(4)"],
         [{"a": 1}, "data.update()"],
     ],
 )
 def test_rejects_calls_with_side_effects(data, bad):
     context = limited(data=data)
 
     with pytest.raises(GuardRejection):
         guarded_eval(bad, context)
 
 
 @pytest.mark.parametrize(
     "code,expected",
     [
         ["(1\n+\n1)", 2],
         ["list(range(10))[-1:]", [9]],
         ["list(range(20))[3:-2:3]", [3, 6, 9, 12, 15]],
     ],
 )
 @pytest.mark.parametrize("context", LIMITED_OR_HIGHER)
 def test_evaluates_complex_cases(code, expected, context):
     assert guarded_eval(code, context()) == expected
 
 
 @pytest.mark.parametrize(
     "code,expected",
     [
         ["1", 1],
         ["1.0", 1.0],
         ["0xdeedbeef", 0xDEEDBEEF],
         ["True", True],
         ["None", None],
         ["{}", {}],
         ["[]", []],
     ],
 )
 @pytest.mark.parametrize("context", MINIMAL_OR_HIGHER)
 def test_evaluates_literals(code, expected, context):
     assert guarded_eval(code, context()) == expected
 
 
 @pytest.mark.parametrize(
     "code,expected",
     [
         ["-5", -5],
         ["+5", +5],
         ["~5", -6],
     ],
 )
 @pytest.mark.parametrize("context", LIMITED_OR_HIGHER)
 def test_evaluates_unary_operations(code, expected, context):
     assert guarded_eval(code, context()) == expected
 
 
 @pytest.mark.parametrize(
     "code,expected",
     [
         ["1 + 1", 2],
         ["3 - 1", 2],
         ["2 * 3", 6],
         ["5 // 2", 2],
         ["5 / 2", 2.5],
         ["5**2", 25],
         ["2 >> 1", 1],
         ["2 << 1", 4],
         ["1 | 2", 3],
         ["1 & 1", 1],
         ["1 & 2", 0],
     ],
 )
 @pytest.mark.parametrize("context", LIMITED_OR_HIGHER)
 def test_evaluates_binary_operations(code, expected, context):
     assert guarded_eval(code, context()) == expected
 
 
 @pytest.mark.parametrize(
     "code,expected",
     [
         ["2 > 1", True],
         ["2 < 1", False],
         ["2 <= 1", False],
         ["2 <= 2", True],
         ["1 >= 2", False],
         ["2 >= 2", True],
         ["2 == 2", True],
         ["1 == 2", False],
         ["1 != 2", True],
         ["1 != 1", False],
         ["1 < 4 < 3", False],
         ["(1 < 4) < 3", True],
         ["4 > 3 > 2 > 1", True],
         ["4 > 3 > 2 > 9", False],
         ["1 < 2 < 3 < 4", True],
         ["9 < 2 < 3 < 4", False],
         ["1 < 2 > 1 > 0 > -1 < 1", True],
         ["1 in [1] in [[1]]", True],
         ["1 in [1] in [[2]]", False],
         ["1 in [1]", True],
         ["0 in [1]", False],
         ["1 not in [1]", False],
         ["0 not in [1]", True],
         ["True is True", True],
         ["False is False", True],
         ["True is False", False],
         ["True is not True", False],
         ["False is not True", True],
     ],
 )
 @pytest.mark.parametrize("context", LIMITED_OR_HIGHER)
 def test_evaluates_comparisons(code, expected, context):
     assert guarded_eval(code, context()) == expected
 
 
 def test_guards_comparisons():
     class GoodEq(int):
         pass
 
     class BadEq(int):
         def __eq__(self, other):
             assert False
 
     context = limited(bad=BadEq(1), good=GoodEq(1))
 
     with pytest.raises(GuardRejection):
         guarded_eval("bad == 1", context)
 
     with pytest.raises(GuardRejection):
         guarded_eval("bad != 1", context)
 
     with pytest.raises(GuardRejection):
         guarded_eval("1 == bad", context)
 
     with pytest.raises(GuardRejection):
         guarded_eval("1 != bad", context)
 
     assert guarded_eval("good == 1", context) is True
     assert guarded_eval("good != 1", context) is False
     assert guarded_eval("1 == good", context) is True
     assert guarded_eval("1 != good", context) is False
 
 
 def test_guards_unary_operations():
     class GoodOp(int):
         pass
 
     class BadOpInv(int):
         def __inv__(self, other):
             assert False
 
     class BadOpInverse(int):
         def __inv__(self, other):
             assert False
 
     context = limited(good=GoodOp(1), bad1=BadOpInv(1), bad2=BadOpInverse(1))
 
     with pytest.raises(GuardRejection):
         guarded_eval("~bad1", context)
 
     with pytest.raises(GuardRejection):
         guarded_eval("~bad2", context)
 
 
 def test_guards_binary_operations():
     class GoodOp(int):
         pass
 
     class BadOp(int):
         def __add__(self, other):
             assert False
 
     context = limited(good=GoodOp(1), bad=BadOp(1))
 
     with pytest.raises(GuardRejection):
         guarded_eval("1 + bad", context)
 
     with pytest.raises(GuardRejection):
         guarded_eval("bad + 1", context)
 
     assert guarded_eval("good + 1", context) == 2
     assert guarded_eval("1 + good", context) == 2
 
 
 def test_guards_attributes():
     class GoodAttr(float):
         pass
 
     class BadAttr1(float):
         def __getattr__(self, key):
             assert False
 
     class BadAttr2(float):
         def __getattribute__(self, key):
             assert False
 
     context = limited(good=GoodAttr(0.5), bad1=BadAttr1(0.5), bad2=BadAttr2(0.5))
 
     with pytest.raises(GuardRejection):
         guarded_eval("bad1.as_integer_ratio", context)
 
     with pytest.raises(GuardRejection):
         guarded_eval("bad2.as_integer_ratio", context)
 
     assert guarded_eval("good.as_integer_ratio()", context) == (1, 2)
 
 
 @pytest.mark.parametrize("context", MINIMAL_OR_HIGHER)
 def test_access_builtins(context):
     assert guarded_eval("round", context()) == round
 
 
 def test_access_builtins_fails():
     context = limited()
     with pytest.raises(NameError):
         guarded_eval("this_is_not_builtin", context)
 
 
 def test_rejects_forbidden():
     context = forbidden()
     with pytest.raises(GuardRejection):
         guarded_eval("1", context)
 
 
 def test_guards_locals_and_globals():
     context = EvaluationContext(
         locals={"local_a": "a"}, globals={"global_b": "b"}, evaluation="minimal"
     )
 
     with pytest.raises(GuardRejection):
         guarded_eval("local_a", context)
 
     with pytest.raises(GuardRejection):
         guarded_eval("global_b", context)
 
 
 def test_access_locals_and_globals():
     context = EvaluationContext(
         locals={"local_a": "a"}, globals={"global_b": "b"}, evaluation="limited"
     )
     assert guarded_eval("local_a", context) == "a"
     assert guarded_eval("global_b", context) == "b"
 
 
 @pytest.mark.parametrize(
     "code",
     ["def func(): pass", "class C: pass", "x = 1", "x += 1", "del x", "import ast"],
 )
 @pytest.mark.parametrize("context", [minimal(), limited(), unsafe()])
 def test_rejects_side_effect_syntax(code, context):
     with pytest.raises(SyntaxError):
         guarded_eval(code, context)
 
 
 def test_subscript():
     context = EvaluationContext(
         locals={}, globals={}, evaluation="limited", in_subscript=True
     )
     empty_slice = slice(None, None, None)
     assert guarded_eval("", context) == tuple()
     assert guarded_eval(":", context) == empty_slice
     assert guarded_eval("1:2:3", context) == slice(1, 2, 3)
     assert guarded_eval(':, "a"', context) == (empty_slice, "a")
 
 
 def test_unbind_method():
     class X(list):
         def index(self, k):
             return "CUSTOM"
 
     x = X()
     assert _unbind_method(x.index) is X.index
     assert _unbind_method([].index) is list.index
     assert _unbind_method(list.index) is None
 
 
 def test_assumption_instance_attr_do_not_matter():
     """This is semi-specified in Python documentation.
 
     However, since the specification says 'not guaranteed
     to work' rather than 'is forbidden to work', future
     versions could invalidate this assumptions. This test
     is meant to catch such a change if it ever comes true.
     """
 
     class T:
         def __getitem__(self, k):
             return "a"
 
         def __getattr__(self, k):
             return "a"
 
     def f(self):
         return "b"
 
     t = T()
     t.__getitem__ = f
     t.__getattr__ = f
     assert t[1] == "a"
     assert t[1] == "a"
 
 
 def test_assumption_named_tuples_share_getitem():
     """Check assumption on named tuples sharing __getitem__"""
     from typing import NamedTuple
 
     class A(NamedTuple):
         pass
 
     class B(NamedTuple):
         pass
 
     assert A.__getitem__ == B.__getitem__
 
 
 @dec.skip_without("numpy")
 def test_module_access():
     import numpy
 
     context = limited(numpy=numpy)
     assert guarded_eval("numpy.linalg.norm", context) == numpy.linalg.norm
 
     context = minimal(numpy=numpy)
     with pytest.raises(GuardRejection):
         guarded_eval("np.linalg.norm", context)