Recursion fix and minor updates

Author: cretz

temporalio/worker/workflow_sandbox/importer.py

@@ -94,13 +94,12 @@ def restrict_built_in(name: str, orig: Any, *args, **kwargs):
                     or k in builtin_matcher.use
                     or k in builtin_matcher.children
                 ):
+                    thread_local = _get_thread_local_builtin(k)
                     self.restricted_builtins.append(
                         (
                             k,
-                            _get_thread_local_builtin(k),
-                            functools.partial(
-                                restrict_built_in, k, getattr(builtins, k)
-                            ),
+                            thread_local,
+                            functools.partial(restrict_built_in, k, thread_local.orig),
                         )
                     )
 
@@ -275,18 +274,18 @@ def applied(self, obj: Any, attr: str, current: _T) -> Iterator[None]:
         # Function carefully crafted to support nesting and situations where
         # other threads may have already set this on obj
         orig_current = self.maybe_current
-        self.current = current
         orig_value = getattr(obj, attr)
         if orig_value is not self:
             setattr(obj, attr, self)
+        self.current = current
         try:
             yield None
         finally:
+            setattr(obj, attr, orig_value)
             if orig_current is None:
                 del self.current
             else:
                 self.current = orig_current
-            setattr(obj, attr, orig_value)
 
     @contextmanager
     def unapplied(self) -> Iterator[None]:

temporalio/worker/workflow_sandbox/runner.py

@@ -33,8 +33,6 @@ def __init__(
         self,
         *,
         restrictions: SandboxRestrictions = SandboxRestrictions.default,
-        # TODO(cretz): Document that this is re-imported and instantiated for
-        # _each_ workflow run.
         runner_class: Type[WorkflowRunner] = UnsandboxedWorkflowRunner,
     ) -> None:
         """Create the sandboxed workflow runner.
@@ -102,7 +100,6 @@ def __init__(
 
         # Create the instance
         self.globals_and_locals = {
-            # "__builtins__": __builtins__.copy(),  # type: ignore
             "__file__": "workflow_sandbox.py",
         }
         self._create_instance()