From 20cf7607bb2bdb1edc8c51a039d3461a09ca7c8e Mon Sep 17 00:00:00 2001
From: Daiyi Peng <daiyip@google.com>
Date: Mon, 13 Jan 2025 13:15:32 -0800
Subject: [PATCH] Expose `lf.coding.permission` and
 `lf.coding.maybe_sandbox_call`.

PiperOrigin-RevId: 715083241
---
 pyglove/core/coding/__init__.py       |  3 +++
 pyglove/core/coding/execution.py      |  6 +++---
 pyglove/core/coding/execution_test.py | 10 +++++-----
 3 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/pyglove/core/coding/__init__.py b/pyglove/core/coding/__init__.py
index b73005f..182631c 100644
--- a/pyglove/core/coding/__init__.py
+++ b/pyglove/core/coding/__init__.py
@@ -22,6 +22,8 @@
 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
 
@@ -29,6 +31,7 @@
 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
index 00fb164..7f36a46 100644
--- 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
index 39d3113..d7c320e 100644
--- 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):