environment bugfix

Author: smacke

nbpipes/macro_tracer.py

@@ -111,15 +111,24 @@ def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.arg_replacer = _ArgReplacer()
         self.lambda_cache: dict[tuple[int, int, TraceEvent], Any] = {}
+        self._overridden_builtins: list[str] = []
         with self.register_additional_ast_bookkeeping():
             self.placeholder_inference_skip_nodes: set[int] = set()
         user_ns = get_user_ns()
         for macro_name, macro in self.macros.items():
             if hasattr(builtins, macro_name):
                 continue
             setattr(builtins, macro_name, macro)
+            self._overridden_builtins.append(macro_name)
             if user_ns is not None:
-                user_ns[macro_name] = macro
+                user_ns.setdefault(macro_name, macro)
+
+    def reset(self) -> None:
+        for macro_name in self._overridden_builtins:
+            if hasattr(builtins, macro_name):
+                delattr(builtins, macro_name)
+        self._overridden_builtins.clear()
+        super().reset()
 
     class _IdentitySubscript:
         def __getitem__(self, item):

nbpipes/pipeline_tracer.py

@@ -234,6 +234,7 @@ def __init__(self, *args, **kwargs) -> None:
         self.binop_arg_nodes_to_skip: set[int] = set()
         self.binop_nodes_to_eval: set[int] = set()
         self.lexical_chain_stack: pyc.TraceStack = self.make_stack()
+        self._overridden_builtins: list[str] = []
         with self.register_additional_ast_bookkeeping():
             self.placeholder_arg_position_cache: dict[int, list[str]] = {}
         self.exc_to_propagate: Exception | None = None
@@ -245,9 +246,17 @@ def __init__(self, *args, **kwargs) -> None:
             extra_builtin_name = extra_builtin.__name__
             if hasattr(builtins, extra_builtin_name):
                 continue
+            self._overridden_builtins.append(extra_builtin_name)
             setattr(builtins, extra_builtin_name, extra_builtin)
-            if user_ns is not None:
-                user_ns[extra_builtin_name] = extra_builtin
+            if user_ns is not None and extra_builtin_name:
+                user_ns.setdefault(extra_builtin_name, extra_builtin)
+
+    def reset(self) -> None:
+        for extra_builtin_name in self._overridden_builtins:
+            if hasattr(builtins, extra_builtin_name):
+                delattr(builtins, extra_builtin_name)
+        self._overridden_builtins.clear()
+        super().reset()
 
     @pyc.register_handler(pyc.before_call, when=is_partial_call, reentrant=True)
     def curry_partial_calls(self, ret, node: ast.Call, *_, **__):
@@ -350,9 +359,8 @@ def should_propagate_handler_exception(
     def maybe_skip_binop_arg(self, ret: object, node_id: int, *_, **__):
         if node_id in self.binop_arg_nodes_to_skip:
             self.binop_arg_nodes_to_skip.remove(node_id)
-            return _skip_binop_args_lambda
-        else:
-            return ret
+            ret = _skip_binop_args_lambda
+        return ret
 
     def reorder_placeholder_names_for_prior_positions(
         self, node: ast.expr, placeholder_names: list[str]

nbpipes/utils.py

@@ -3,7 +3,8 @@
 from typing import Any
 
 
-def get_user_ns() -> dict[str, Any] | None:
+# for unittest.mock patching
+def _get_user_ns_impl() -> dict[str, Any] | None:
     try:
         from IPython import get_ipython
 
@@ -13,3 +14,7 @@ def get_user_ns() -> dict[str, Any] | None:
     except ImportError:
         pass
     return None
+
+
+def get_user_ns() -> dict[str, Any] | None:
+    return _get_user_ns_impl()