fix(mirage-sandbox): don't re-home imported functions onto the sandbox namespace

synalinks/src/sandboxes/mirage_sandbox.py

@@ -661,9 +661,22 @@ class _SockFprog(ctypes.Structure):
 # leaves ``f`` raising NameError forever), and every dump/restore round-trip
 # nests another ghost copy into the state file. Rebuilding each function on
 # ``ns`` restores true REPL semantics: one shared global namespace.
+#
+# Only functions *defined in the sandbox* may be re-homed. An imported
+# function closes over its own module's globals, and rebuilding it on ``ns``
+# strips the module internals it reaches at call time: re-homing
+# ``collections.Counter``'s methods makes ``Counter('aa')`` raise
+# ``NameError: name '_collections_abc' is not defined``, and re-homing
+# ``os.path.join`` loses ``sep``. Sandbox-defined functions are the ones whose
+# (ghost) globals are a copy of ``ns``, which carries ``__name__ ==
+# "__main__"``; an imported one carries its own module name.
 import types as _types
 def _rehome(value):
-    if isinstance(value, _types.FunctionType) and value.__globals__ is not ns:
+    if (
+        isinstance(value, _types.FunctionType)
+        and value.__globals__ is not ns
+        and value.__globals__.get("__name__") == "__main__"
+    ):
         fixed = _types.FunctionType(
             value.__code__, ns, value.__name__, value.__defaults__, value.__closure__
         )

synalinks/src/sandboxes/mirage_sandbox_test.py

@@ -116,6 +116,35 @@ async def test_rehomed_function_keeps_closure_and_kwdefaults(self):
         result = await sandbox.run("print(add5(1), add5(1, scale=1))")
         self.assertIn("12 6", result.stdout)
 
+    async def test_imported_function_keeps_its_own_module_globals(self):
+        """Re-homing must not touch functions imported from a module.
+
+        An imported function closes over its *module's* globals and reaches
+        names in them at call time. Rebuilding it on the sandbox namespace
+        strips those: ``os.path.join`` re-homed this way loses ``sep``."""
+        sandbox = MirageSandbox(timeout=_TIMEOUT)
+        await sandbox.run("from os.path import join")
+        result = await sandbox.run("print(join('a', 'b'))")
+        self.assertIn("a/b", result.stdout)
+        self.assertFalse(result.error)
+
+    async def test_imported_class_methods_keep_their_module_globals(self):
+        """Same for the methods of an imported class.
+
+        ``Counter.update`` reads ``_collections_abc`` and ``_count_elements``
+        from ``collections``' own globals, so re-homing it onto the sandbox
+        namespace made ``Counter('aa')`` raise NameError on every run after
+        the one that imported it. Worse, the class is re-homed *in place*, so
+        a later plain ``import collections`` inherited the damage."""
+        sandbox = MirageSandbox(timeout=_TIMEOUT)
+        await sandbox.run("from collections import Counter")
+        result = await sandbox.run("print(Counter('aab')['a'])")
+        self.assertIn("2", result.stdout)
+        self.assertFalse(result.error)
+        fresh = await sandbox.run("import collections\nprint(collections.Counter('aab')['a'])")
+        self.assertIn("2", fresh.stdout)
+        self.assertFalse(fresh.error)
+
     async def test_state_survives_error(self):
         sandbox = MirageSandbox(timeout=_TIMEOUT)
         await sandbox.run("keep = 11")