pyglove 0.4.5.dev202501060809__py3-none-any.whl → 0.4.5.dev202501110808__py3-none-any.whl

pyglove/core/__init__.py

@@ -302,6 +302,9 @@ docstr = utils.docstr
 catch_errors = utils.catch_errors
 timeit = utils.timeit
 
+colored = utils.colored
+decolor = utils.decolor
+
 # Symbols from 'views' sub-module.
 
 from pyglove.core import views

pyglove/core/coding/__init__.py

@@ -18,6 +18,19 @@
 # pylint: disable=g-bad-import-order
 # pylint: disable=g-importing-member
 
+from pyglove.core.coding.errors import CodeError
+from pyglove.core.coding.errors import SerializationError
+
+from pyglove.core.coding.permissions import CodePermission
+
+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 run
+
 from pyglove.core.coding.function_generation import NO_TYPE_ANNOTATION
 from pyglove.core.coding.function_generation import make_function
 

pyglove/core/coding/errors.py

@@ -0,0 +1,111 @@
+# Copyright 2025 The PyGlove Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Python code errors."""
+
+import io
+import sys
+import textwrap
+import traceback
+from typing import Optional
+
+from pyglove.core import utils
+
+
+class CodeError(RuntimeError):
+  """Python code error."""
+
+  def __init__(
+      self,
+      code: str,
+      cause: Exception,
+      ):
+    self.code = code
+    self.cause = cause
+
+    # Figure out the starting and ending line numbers of the erratic code.
+    lineno = None
+    end_lineno = None
+    if isinstance(cause, SyntaxError):
+      lineno = cause.lineno
+      # For Python 3.9 and below, `end_lineno` is not available.
+      end_lineno = getattr(cause, 'end_lineno', lineno)
+    elif not isinstance(cause, TimeoutError):
+      tb = sys.exc_info()[2]
+      frames = traceback.extract_tb(tb, limit=5)
+      for f in frames:
+        if not f.filename or f.filename == '<string>':
+          lineno = f.lineno
+          end_lineno = lineno
+          break
+    self.lineno = lineno
+    self.end_lineno = end_lineno
+
+  def __str__(self):
+    return self.format(include_complete_code=True)
+
+  def code_lines(self, start_line: int, end_line: int):
+    """Returns code lines ."""
+    return '\n'.join(self.code.split('\n')[start_line:end_line])
+
+  def format(self, include_complete_code: bool = True):
+    """Formats the code error."""
+    r = io.StringIO()
+    error_message = str(self.cause).rstrip()
+    if 'line' not in error_message and self.lineno is not None:
+      error_message += f' (<unknown>, line {self.lineno})'
+    r.write(
+        utils.colored(
+            f'{self.cause.__class__.__name__}: {error_message}', 'magenta'))
+
+    if self.lineno is not None:
+      r.write('\n\n')
+      r.write(textwrap.indent(
+          utils.colored(
+              self.code_lines(self.lineno - 1, self.end_lineno), 'magenta'),
+          ' ' * 2
+      ))
+      r.write('\n')
+
+    if include_complete_code:
+      r.write('\n')
+      r.write(utils.colored('[Code]', 'green', styles=['bold']))
+      r.write('\n\n')
+      r.write(utils.colored('  ```python\n', 'green'))
+      r.write(textwrap.indent(
+          utils.colored(self.code, 'green'),
+          ' ' * 2
+      ))
+      r.write(utils.colored('\n  ```\n', 'green'))
+    return r.getvalue()
+
+
+class SerializationError(RuntimeError):
+  """Object serialization error."""
+
+  def __init__(self, message: Optional[str], cause: Exception):
+    self.message = message
+    self.cause = cause
+
+  def __str__(self):
+    r = io.StringIO()
+    cause_message = str(self.cause).rstrip()
+    if self.message:
+      r.write(utils.colored(self.message, 'magenta'))
+      r.write('\n\n')
+    r.write(
+        utils.colored(
+            f'{self.cause.__class__.__name__}: {cause_message}', 'magenta'
+        )
+    )
+    return r.getvalue()

pyglove/core/coding/errors_test.py

@@ -0,0 +1,98 @@
+# Copyright 2025 The PyGlove Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import inspect
+import unittest
+
+from pyglove.core.coding import errors
+from pyglove.core.coding import execution
+
+
+def code_error(code: str) -> errors.CodeError:
+  try:
+    execution.run(inspect.cleandoc(code), timeout=2)
+    assert False, 'should not reach here'
+  except errors.CodeError as e:
+    return e
+
+
+class CodeErrorsTest(unittest.TestCase):
+
+  def test_format(self):
+    e = code_error(
+        """
+        x = y + 1
+        """
+    )
+    self.assertIn('[Code]', str(e))
+    self.assertNotIn(
+        '[Code]', e.format(include_complete_code=False))
+
+  def test_lineno(self):
+    self.assertEqual(
+        code_error(
+            """
+            x = y + 1
+            """
+        ).lineno, 1)
+    self.assertEqual(
+        code_error(
+            """
+            x = 1
+            for i of x:
+              y = i
+            """
+        ).lineno, 2)
+    self.assertEqual(
+        code_error(
+            """
+            x = 1
+            y = 2
+            raise ValueError
+            """
+        ).lineno, 3)
+
+  def test_lineno_in_error_message(self):
+    def assert_lineno(code):
+      e = code_error(code)
+      self.assertIn('line', e.format(include_complete_code=False))
+
+    assert_lineno(
+        """
+        x = y + 1
+        """
+    )
+    assert_lineno(
+        """
+        x = 1
+          y = 2
+        """
+    )
+    assert_lineno(
+        """
+        raise ValueError()
+        """
+    )
+
+
+class SerializationErrorTest(unittest.TestCase):
+
+  def test_str(self):
+    e = errors.SerializationError(
+        'Output cannot be serialized.', ValueError('abc'))
+    self.assertIn('Output cannot be serialized', str(e))
+    self.assertIn('ValueError: abc', str(e))
+
+
+if __name__ == '__main__':
+  unittest.main()

pyglove/core/coding/execution.py

@@ -0,0 +1,309 @@
+# Copyright 2025 The PyGlove Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Python code execution."""
+
+import ast
+import contextlib
+import io
+import multiprocessing
+import pickle
+from typing import Any, Callable, Dict, Optional, Union
+
+from pyglove.core import utils
+from pyglove.core.coding import errors
+from pyglove.core.coding import parsing
+from pyglove.core.coding import permissions
+
+
+# Key in returned dict that captures stdout.
+STDOUT_KEY = '__stdout__'
+
+# Key in the returned dict that represents the final result.
+RESULT_KEY = '__result__'
+_TLS_CODE_RUN_CONTEXT = '__code_run_context__'
+
+
+@contextlib.contextmanager
+def context(**kwargs):
+  """Context manager to inject symbols for code execution."""
+  ctx = get_context()
+  ctx.update(kwargs)
+  utils.thread_local_push(_TLS_CODE_RUN_CONTEXT, ctx)
+
+  try:
+    yield ctx
+  finally:
+    utils.thread_local_pop(_TLS_CODE_RUN_CONTEXT)
+
+
+def get_context() -> Dict[str, Any]:
+  """Gets the current context for code execution."""
+  context_stack = utils.thread_local_get(_TLS_CODE_RUN_CONTEXT, None)
+  return dict(context_stack[-1]) if context_stack else {}
+
+
+def evaluate(
+    code: str,
+    *,
+    global_vars: Optional[Dict[str, Any]] = None,
+    permission: Optional[permissions.CodePermission] = None,
+    returns_stdout: bool = False,
+    outputs_intermediate: bool = False,
+) -> Union[Any, Dict[str, Any]]:
+  """Executes Python code.
+
+  Features:
+    * Fine-grained execution policy for limiting what APIs could be executed.
+      This eliminates the need for sandboxing.
+    * It exposes both the final results and intermediate results (variables).
+
+  Args:
+    code: Python code to run.
+    global_vars: An optional dict as the globals that could be referenced by the
+      code.
+    permission: Permission for the Python code to run.
+    returns_stdout: If True, the stdout (a str) will be returned.
+    outputs_intermediate: Applicable when returns_stdout is False. If True,
+      intermediate output will be outputted as a dict, with the last line's
+      value accessible by key '__result__' and the std output accessible by
+      key '__stdout__'. Otherwise the value of the last line will be returned.
+
+  Returns:
+    The value of the last line of the code block. Or a dict of variable
+    names of all locals to their evaluated values as the output of the code to
+    run. The value for the last line can be accessed by key '__result__'. Or the
+    stdout as a str.
+  """
+  # Set up the permission and context.
+  permission = permission or permissions.get_permission()
+  ctx = dict(get_context())
+  if global_vars:
+    ctx.update(global_vars)
+
+  # Parse the code str.
+  code_block = parsing.parse(code, permission)
+  global_vars, orig_global_vars = ctx, ctx.copy()
+
+  # No code.
+  if not code_block.body:   # pytype: disable=attribute-error
+    return {} if outputs_intermediate else None
+
+  stdout = io.StringIO()
+  with contextlib.redirect_stdout(stdout):
+    if hasattr(code_block.body[-1], 'value'):   # pytype: disable=attribute-error
+      last_expr = code_block.body.pop()  # pytype: disable=attribute-error
+      result_vars = [RESULT_KEY]
+
+      if isinstance(last_expr, ast.Assign):
+        for name_node in last_expr.targets:
+          result_vars.append(name_node.id)
+
+      last_expr = ast.Expression(last_expr.value)  # pytype: disable=attribute-error
+
+      try:
+        # Execute the lines before the last expression.
+        # NOTE(daiyip): Only a `globals` dict is specified here, which will also
+        # be used to output intermediate values by `exec`. We do not specify a
+        # separate `locals` dict here, for - "If exec gets two separate objects
+        # as globals and locals, the code will be executed as if it were
+        # embedded in a class definition." - as the Python document explains.
+        # The outcome is that new functions defined in the code block could not
+        # be called by other newly defined functions.
+        # Refer to https://stackoverflow.com/questions/
+        # 73940751/why-cant-i-call-a-function-from-another-function-using-exec
+        # for more details.
+        exec(compile(code_block, '', mode='exec'), global_vars)  # pylint: disable=exec-used
+
+        # Evaluate the last expression.
+        result = eval(  # pylint: disable=eval-used
+            compile(last_expr, '', mode='eval'), global_vars
+        )
+      except Exception as e:
+        raise errors.CodeError(code, e) from e
+
+      for result_var in result_vars:
+        global_vars[result_var] = result
+    else:
+      try:
+        exec(compile(code_block, '', mode='exec'), global_vars)  # pylint: disable=exec-used
+      except Exception as e:
+        raise errors.CodeError(code, e) from e
+      global_vars[RESULT_KEY] = list(global_vars.values())[-1]
+
+  if returns_stdout:
+    return stdout.getvalue()
+  if outputs_intermediate:
+    outputs = {}
+    for k, v in global_vars.items():
+      if k == '__builtins__':
+        continue
+      if k not in orig_global_vars or v is not orig_global_vars[k]:
+        outputs[k] = v
+    # Add stdout to outputs.
+    outputs[STDOUT_KEY] = stdout.getvalue()
+    return outputs
+  return global_vars[RESULT_KEY]
+
+
+def sandbox_call(
+    func: Callable[..., Any],
+    *args,
+    timeout: Optional[float] = None,
+    **kwargs) -> Any:
+  """Calls a function with sandboxing.
+
+  Args:
+    func: Function to call.
+    *args: Positional arguments for `func`
+    timeout: Execution timeout in seconds. If None, wait `func` to complete.
+    **kwargs: Keyword arguments for `func`.
+
+  Returns:
+    Return value from `func`.
+
+  Raises:
+    TimeoutError: If the execution time exceeds the timeout.
+    Exception: Exception raised from `func`.
+  """
+  def _call(q, *args, **kwargs):
+    # NOTE(daiyip): if `q` is closed by the main process when `q.put` is called
+    # on a subprocess, ValueError will be raised. This is okay since the main
+    # process is no longer waiting for the result, and the subprocess could
+    # recycled with non-zero error code, which does not affect the main
+    # process.
+    def _run():
+      r = func(*args, **kwargs)
+      try:
+        return pickle.dumps(r)
+      except Exception as e:
+        raise errors.SerializationError(
+            f'Cannot serialize sandbox result: {r}', e
+        ) from e
+
+    try:
+      q.put(_run())
+    except Exception as e:  # pylint: disable=broad-exception-caught
+      q.put(e)
+
+  q = multiprocessing.Queue()
+  try:
+    p = multiprocessing.Process(
+        target=_call, args=tuple([q] + list(args)), kwargs=kwargs)
+    p.start()
+    p.join(timeout=timeout)
+    if p.is_alive():
+      # We use `kill` instead of `terminate` to release process resources
+      # right away.
+      p.kill()
+      raise TimeoutError(f'Execution time exceed {timeout} seconds.')
+    x = q.get()
+    if isinstance(x, Exception):
+      raise x
+    try:
+      return pickle.loads(x)
+    except Exception as e:
+      raise errors.SerializationError(
+          'Cannot deserialize the output from sandbox.', e
+      ) from e
+  finally:
+    q.close()
+
+
+def _maybe_call_in_sandbox(
+    func: Callable[..., Any],
+    *args,
+    sandbox: Optional[bool] = None,
+    timeout: Optional[float] = None,
+    **kwargs
+) -> Any:
+  """Calls a function with sandbox support.
+
+  Args:
+    func: Function to call.
+    *args: Postional args that will be passed to `func`.
+    sandbox: If True, run code in sandbox; If False, run code in current
+      process. If None, run in sandbox first, if the output could not be
+      serialized and pass to current process, run the code again in current
+      process.
+    timeout: Execution timeout in seconds. If None, wait the code the complete.
+    **kwargs: Keyword args that will be passed to `func`.
+
+  Returns:
+    The return value of `func`.
+
+  Raises:
+    TimeoutError: If the execution time exceeds the timeout.
+    Exception: Exception  that are raised from `func`.
+  """
+  if sandbox is None:
+    try:
+      return sandbox_call(func, *args, timeout=timeout, **kwargs)
+    # NOTE(daiyip): output could be serialized across processes, giving it
+    # already finishes on sandbox, so it should be much safer to run under
+    # current process.
+    except errors.SerializationError:
+      return func(*args, **kwargs)
+  elif sandbox:
+    return sandbox_call(func, *args, timeout=timeout, **kwargs)
+  else:
+    return func(*args, **kwargs)
+
+
+def run(
+    code: str,
+    *,
+    global_vars: Optional[Dict[str, Any]] = None,
+    permission: Optional[permissions.CodePermission] = None,
+    returns_stdout: bool = False,
+    outputs_intermediate: bool = False,
+    sandbox: Optional[bool] = None,
+    timeout: Optional[float] = None,
+) -> Union[Any, Dict[str, Any]]:
+  """Executes Python code.
+
+  Features:
+    * Fine-grained execution policy for limiting what APIs could be executed.
+      This eliminates the need for sandboxing.
+    * It exposes both the final results and intermediate results (variables).
+
+  Args:
+    code: Python code to run.
+    global_vars: An optional dict of
+    permission: Permission for the Python code to run.
+    returns_stdout: If True, the stdout (a str) will be returned.
+    outputs_intermediate: Applicable when returns_stdout is False. If True,
+      intermediate output will be outputted as a dict, with the last line's
+      value accessible by key '__result__' and the std output accessible by
+      key '__stdout__'. Otherwise the value of the last line will be returned.
+    sandbox: If True, run code in sandbox; If False, run code in current
+      process. If None, run in sandbox first, if the output could not be
+      serialized and pass to current process, run the code again in current
+      process.
+    timeout: Execution timeout in seconds. If None, wait the code the complete.
+
+  Returns:
+    The value of the last line of the code block. Or a dict of variable
+    names of all locals to their evaluated values as the output of the code to
+    run. The value for the last line can be accessed by key '__result__'. Or the
+    stdout as a str.
+
+  Raises:
+    TimeoutError: If the execution time exceeds the timeout.
+    Exception: Exception  that are raised from the code.
+  """
+  return _maybe_call_in_sandbox(
+      evaluate, code=code, global_vars=global_vars, permission=permission,
+      returns_stdout=returns_stdout, outputs_intermediate=outputs_intermediate,
+      sandbox=sandbox, timeout=timeout
+  )