Expose lf.coding.permission and lf.coding.maybe_sandbox_call.

PiperOrigin-RevId: 715083241
google · Jan 13, 2025 · 20cf760 · 20cf760
1 parent faca119
commit 20cf760
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 8 deletions.
diff --git a/pyglove/core/coding/__init__.py b/pyglove/core/coding/__init__.py
@@ -22,13 +22,16 @@
 from pyglove.core.coding.errors import SerializationError
 
 from pyglove.core.coding.permissions import CodePermission
+from pyglove.core.coding.permissions import permission
+from pyglove.core.coding.permissions import get_permission
 
 from pyglove.core.coding.parsing import parse
 
 from pyglove.core.coding.execution import context
 from pyglove.core.coding.execution import get_context
 from pyglove.core.coding.execution import evaluate
 from pyglove.core.coding.execution import sandbox_call
+from pyglove.core.coding.execution import maybe_sandbox_call
 from pyglove.core.coding.execution import run
 
 from pyglove.core.coding.function_generation import NO_TYPE_ANNOTATION

diff --git a/pyglove/core/coding/execution.py b/pyglove/core/coding/execution.py
@@ -220,14 +220,14 @@ def _run():
     q.close()
 
 
-def _maybe_call_in_sandbox(
+def maybe_sandbox_call(
     func: Callable[..., Any],
     *args,
     sandbox: Optional[bool] = None,
     timeout: Optional[float] = None,
     **kwargs
 ) -> Any:
-  """Calls a function with sandbox support.
+  """Maybe calls a function with sandboxing.
 
   Args:
     func: Function to call.
@@ -302,7 +302,7 @@ def run(
     TimeoutError: If the execution time exceeds the timeout.
     Exception: Exception  that are raised from the code.
   """
-  return _maybe_call_in_sandbox(
+  return maybe_sandbox_call(
       evaluate, code=code, global_vars=global_vars, permission=permission,
       returns_stdout=returns_stdout, outputs_intermediate=outputs_intermediate,
       sandbox=sandbox, timeout=timeout

diff --git a/pyglove/core/coding/execution_test.py b/pyglove/core/coding/execution_test.py
@@ -242,7 +242,7 @@ def foo(x, y):
       return x + y
 
     self.assertEqual(
-        execution._maybe_call_in_sandbox(foo, 1, y=2, sandbox=False),
+        execution.maybe_sandbox_call(foo, 1, y=2, sandbox=False),
         3
     )
 
@@ -251,7 +251,7 @@ def foo(x, y):
       return x + y
 
     self.assertEqual(
-        execution._maybe_call_in_sandbox(foo, 1, y=2, sandbox=True),
+        execution.maybe_sandbox_call(foo, 1, y=2, sandbox=True),
         3
     )
 
@@ -262,14 +262,14 @@ class A:
       return A
 
     with self.assertRaises(errors.SerializationError):
-      execution._maybe_call_in_sandbox(make_cls, sandbox=True)
+      execution.maybe_sandbox_call(make_cls, sandbox=True)
 
   def test_call_with_automatic_sandboxing(self):
     def foo(x, y):
       return x + y
 
     self.assertEqual(
-        execution._maybe_call_in_sandbox(foo, 1, y=2),
+        execution.maybe_sandbox_call(foo, 1, y=2),
         3
     )
 
@@ -279,7 +279,7 @@ class A:
         x: str
       return A
 
-    self.assertTrue(inspect.isclass(execution._maybe_call_in_sandbox(make_cls)))
+    self.assertTrue(inspect.isclass(execution.maybe_sandbox_call(make_cls)))
 
 
 class RunTest(unittest.TestCase):