pyglove 0.4.5.dev20240318__py3-none-any.whl → 0.4.5.dev202501132210__py3-none-any.whl

pyglove/core/coding/execution_test.py

@@ -0,0 +1,333 @@
+# 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 dataclasses
+import inspect
+import time
+import unittest
+from pyglove.core.coding import errors
+from pyglove.core.coding import execution
+from pyglove.core.coding import permissions
+
+
+class EvaluateTest(unittest.TestCase):
+
+  def test_with_context(self):
+    with execution.context(x=1, y=0):
+      with execution.context(x=2, z=2):
+        self.assertEqual(
+            execution.evaluate(
+                inspect.cleandoc(
+                    """
+                    p = x + y + z
+                    """
+                ),
+                # Override value from the context.
+                global_vars=dict(z=3),
+                outputs_intermediate=True,
+            ),
+            dict(p=2 + 0 + 3, __result__=2 + 0 + 3, __stdout__=''),
+        )
+
+  def test_basics(self):
+    self.assertEqual(
+        execution.evaluate(
+            inspect.cleandoc(
+                """
+                x = 1
+                y = x + 1
+                print(y)
+                z = x + y
+                """
+            ),
+            outputs_intermediate=True,
+        ),
+        dict(x=1, y=2, z=3, __result__=3, __stdout__='2\n'),
+    )
+    self.assertEqual(
+        execution.evaluate(
+            inspect.cleandoc(
+                """
+                x = 1
+                y = x + 1
+                print(y)
+                z = x + y
+                """
+            ),
+        ),
+        3,
+    )
+    with self.assertRaisesRegex(errors.CodeError, 'ValueError'):
+      execution.evaluate(
+          inspect.cleandoc(
+              """
+              def foo():
+                raise ValueError("intentional error")
+              foo()
+              """
+          ),
+          permission=permissions.CodePermission.ALL
+      )
+
+  def test_class_def(self):
+    ret = execution.evaluate(
+        inspect.cleandoc(
+            """
+            @dataclasses.dataclass
+            class A:
+              x: int
+              y: int
+              def __call__(self):
+                return self.x + self.y
+            """
+        ),
+        permission=permissions.CodePermission.ALL,
+        global_vars=dict(dataclasses=dataclasses),
+        outputs_intermediate=True,
+    )
+    self.assertEqual(list(ret.keys()), ['A', '__result__', '__stdout__'])
+    self.assertTrue(inspect.isclass(ret['A']))
+    self.assertIs(ret['__result__'], ret['A'])
+    self.assertEqual(ret['__stdout__'], '')
+
+  def test_function_def(self):
+    ret = execution.evaluate(
+        inspect.cleandoc(
+            """
+            def foo(x, y):
+              return x + y
+
+            def bar(z):
+              return z + foo(z, z)
+            """
+        ),
+        permission=permissions.CodePermission.ALL,
+        outputs_intermediate=True,
+    )
+    self.assertEqual(
+        list(ret.keys()), ['foo', 'bar', '__result__', '__stdout__']
+    )
+    self.assertTrue(inspect.isfunction(ret['foo']))
+    self.assertTrue(inspect.isfunction(ret['bar']))
+    self.assertIs(ret['__result__'], ret['bar'])
+
+  def test_function_def_and_call(self):
+    code = inspect.cleandoc(
+        """
+        def foo(x, y):
+          return x + y
+
+        def bar(z):
+          print(f'z is {z}')
+          return z + foo(z, z)
+
+        bar(1)
+        """
+    )
+    ret = execution.evaluate(
+        code,
+        permission=permissions.CodePermission.ALL,
+        outputs_intermediate=True,
+    )
+    self.assertEqual(
+        list(ret.keys()), ['foo', 'bar', '__result__', '__stdout__']
+    )
+    self.assertEqual(ret['__result__'], 3)
+    ret = execution.evaluate(
+        code,
+        permission=permissions.CodePermission.ALL,
+        returns_stdout=True,
+    )
+    self.assertEqual(ret, 'z is 1\n')
+
+  def test_complex(self):
+    ret = execution.evaluate(
+        inspect.cleandoc(
+            """
+            @dataclasses.dataclass
+            class A:
+              x: int
+              y: int
+              def __call__(self, z):
+                return self.x + self.y + z
+
+            def foo(x, y):
+              return x + y
+            k = A(1, 2)
+            k(foo(3, 4))
+            """
+        ),
+        permission=permissions.CodePermission.ALL,
+        global_vars=dict(dataclasses=dataclasses),
+        outputs_intermediate=True,
+    )
+    self.assertEqual(
+        list(ret.keys()), ['A', 'foo', 'k', '__result__', '__stdout__']
+    )
+    self.assertTrue(inspect.isclass(ret['A']))
+    self.assertTrue(inspect.isfunction(ret['foo']))
+    self.assertIsInstance(ret['k'], ret['A'])
+    self.assertEqual(ret['__result__'], 10)
+
+  def test_run_with_error(self):
+    with self.assertRaisesRegex(
+        errors.CodeError, 'NameError: name .* is not defined'
+    ):
+      execution.evaluate(
+          inspect.cleandoc(
+              """
+              x = 1
+              y = x + z
+              """
+          ),
+          permission=permissions.CodePermission.ALL,
+      )
+    with self.assertRaisesRegex(errors.CodeError, 'ValueError'):
+      execution.evaluate(
+          'raise ValueError()', permission=permissions.CodePermission.ALL
+      )
+
+
+@dataclasses.dataclass
+class Foo:
+  x: int
+  y: int
+
+
+class SandboxCallTest(unittest.TestCase):
+
+  def test_basics(self):
+    def f(x, y):
+      return x + y
+    self.assertEqual(execution.sandbox_call(f, 1, y=2), 3)
+
+  def test_complex_type(self):
+    def f(x, y):
+      return Foo(x, y)
+
+    self.assertEqual(execution.sandbox_call(f, 1, 2), Foo(1, 2))
+
+  def test_timeout(self):
+    def f(x):
+      time.sleep(x)
+
+    self.assertIsNone(execution.sandbox_call(f, 0, timeout=1))
+    with self.assertRaises(TimeoutError):
+      execution.sandbox_call(f, 2, timeout=1)
+
+  def test_raise(self):
+    def f(x):
+      if x == 0:
+        raise ValueError()
+
+    self.assertIsNone(execution.sandbox_call(f, 1))
+    with self.assertRaises(ValueError):
+      execution.sandbox_call(f, 0)
+
+
+class CallTest(unittest.TestCase):
+
+  def test_call_without_sandboxing(self):
+    def foo(x, y):
+      return x + y
+
+    self.assertEqual(
+        execution.maybe_sandbox_call(foo, 1, y=2, sandbox=False),
+        3
+    )
+
+  def test_call_with_sandboxing(self):
+    def foo(x, y):
+      return x + y
+
+    self.assertEqual(
+        execution.maybe_sandbox_call(foo, 1, y=2, sandbox=True),
+        3
+    )
+
+    def make_cls():
+      @dataclasses.dataclass
+      class A:
+        x: str
+      return A
+
+    with self.assertRaises(errors.SerializationError):
+      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_sandbox_call(foo, 1, y=2),
+        3
+    )
+
+    def make_cls():
+      @dataclasses.dataclass
+      class A:
+        x: str
+      return A
+
+    self.assertTrue(inspect.isclass(execution.maybe_sandbox_call(make_cls)))
+
+
+class RunTest(unittest.TestCase):
+
+  def test_run_without_sandboxing(self):
+    self.assertEqual(
+        execution.run(
+            'x + y',
+            global_vars=dict(x=1, y=2),
+            sandbox=False,
+        ),
+        3,
+    )
+
+  def test_run_with_sandboxing(self):
+    self.assertEqual(
+        execution.run(
+            'x + y',
+            global_vars=dict(x=1, y=2),
+            sandbox=True,
+        ),
+        3,
+    )
+
+  def test_run_with_automatic_sandboxing(self):
+    self.assertEqual(
+        execution.run(
+            'x + y',
+            global_vars=dict(x=1, y=2),
+        ),
+        3,
+    )
+
+    r = execution.run(
+        inspect.cleandoc("""
+            def foo(x, y):
+              return x + y
+
+            @dataclasses.dataclass
+            class A:
+              x: str
+            """),
+        global_vars=dict(dataclasses=dataclasses),
+        outputs_intermediate=True,
+    )
+    self.assertTrue(inspect.isfunction(r['foo']))
+    self.assertTrue(inspect.isclass(r['A']))
+
+
+if __name__ == '__main__':
+  unittest.main()

pyglove/core/{object_utils/codegen.py → coding/function_generation.py}

@@ -11,10 +11,16 @@
 # 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.
-"""Utilities for code generation."""
+"""Utilities for function generation."""
 
 from typing import Any, Dict, List, Optional
-from pyglove.core.object_utils.missing import MISSING_VALUE
+
+
+class _NoTypeAnnotation:
+  """Placeholder for no type annotation."""
+
+
+NO_TYPE_ANNOTATION = _NoTypeAnnotation()
 
 
 def make_function(
@@ -24,11 +30,11 @@ def make_function(
     *,
     exec_globals: Optional[Dict[str, Any]] = None,
     exec_locals: Optional[Dict[str, Any]] = None,
-    return_type: Any = MISSING_VALUE):
+    return_type: Any = NO_TYPE_ANNOTATION):
   """Creates a function dynamically from source."""
   if exec_locals is None:
     exec_locals = {}
-  if return_type != MISSING_VALUE:
+  if return_type != NO_TYPE_ANNOTATION:
     exec_locals['_return_type'] = return_type
     return_annotation = '->_return_type'
   else:

pyglove/core/{object_utils/codegen_test.py → coding/function_generation_test.py}

@@ -1,4 +1,4 @@
-# Copyright 2022 The PyGlove Authors
+# Copyright 2024 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.
@@ -11,20 +11,18 @@
 # 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.
-"""Tests for pyglove.object_utils.codegen."""
-
 import inspect
 import typing
 import unittest
 
-from pyglove.core.object_utils import codegen
+from pyglove.core.coding import function_generation
 
 
 class MakeFunctionTest(unittest.TestCase):
-  """Tests for codegen.make_function."""
+  """Tests for function_generation.make_function."""
 
   def test_make_function_with_type_annotations(self):
-    func = codegen.make_function(
+    func = function_generation.make_function(
         'foo',
         ['x: typing.Optional[int]', 'y: int = 0'],
         ['return x + y'],
@@ -41,7 +39,7 @@ class MakeFunctionTest(unittest.TestCase):
     self.assertEqual(func(1, 2), 3)
 
   def test_make_function_without_type_annotations(self):
-    func = codegen.make_function(
+    func = function_generation.make_function(
         'foo',
         ['x', 'y'],
         ['return x + y'])

pyglove/core/coding/parsing.py

@@ -0,0 +1,153 @@
+# 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 parser with permission control."""
+
+import ast
+import sys
+from typing import Optional
+
+from pyglove.core.coding import errors
+from pyglove.core.coding import permissions
+
+
+class _CodeValidator(ast.NodeVisitor):
+  """Python AST node visitor for ensuring code are permitted."""
+
+  def __init__(
+      self,
+      code: str,
+      permission: permissions.CodePermission
+  ):
+    super().__init__()
+    self.code = code
+    self.permission = permission
+
+  def verify(
+      self,
+      node,
+      flag: permissions.CodePermission,
+      node_type,
+      error_message: str,
+  ) -> None:
+    """Verifies if the node type is permitted based on flag."""
+    if isinstance(node_type, (tuple, list)):
+      node_type = tuple(t for t in node_type if t is not None)
+    if isinstance(node, node_type) and not (self.permission & flag):
+      if sys.version_info >= (3, 10):
+        error = SyntaxError(
+            error_message,
+            (
+                '<generated-code>',
+                node.lineno,
+                node.col_offset,
+                self._code_line(node.lineno),
+                node.end_lineno,
+                node.end_col_offset,
+            )
+        )
+      else:
+        error = SyntaxError(
+            error_message,
+            (
+                '<generated-code>',
+                node.lineno,
+                node.col_offset,
+                self._code_line(node.lineno),
+            )
+        )
+        setattr(error, 'end_lineno', node.lineno)
+        setattr(error, 'end_col_offset', node.col_offset)
+      raise error
+
+  def _code_line(self, lineno):
+    return self.code.split('\n')[lineno - 1]
+
+  def generic_visit(self, node):
+    self.verify(
+        node,
+        permissions.CodePermission.ASSIGN,
+        (ast.Assign),
+        'Assignment is not allowed.',
+    )
+
+    self.verify(
+        node,
+        permissions.CodePermission.CONDITION,
+        # Match is not supported until Python 3.10.
+        (ast.If, getattr(ast, 'Match', None)),
+        'Condition is not allowed.',
+    )
+
+    self.verify(
+        node,
+        permissions.CodePermission.LOOP,
+        (ast.For, ast.While, ast.AsyncFor, ast.AsyncWith),
+        'Loop is not allowed.',
+    )
+
+    self.verify(
+        node,
+        permissions.CodePermission.EXCEPTION,
+        (ast.Try, ast.Raise, ast.Assert),
+        'Exception is not allowed.',
+    )
+
+    self.verify(
+        node,
+        permissions.CodePermission.CALL,
+        ast.Call,
+        'Call is not allowed.',
+    )
+    self.verify(
+        node,
+        permissions.CodePermission.CLASS_DEFINITION,
+        ast.ClassDef,
+        'Class definition is not allowed.',
+    )
+
+    self.verify(
+        node,
+        permissions.CodePermission.FUNCTION_DEFINITION,
+        (
+            ast.FunctionDef,
+            ast.AsyncFunctionDef,
+            ast.Lambda,
+            ast.Return,
+            ast.Yield,
+            ast.YieldFrom,
+        ),
+        'Function definition is not allowed.',
+    )
+
+    self.verify(
+        node,
+        permissions.CodePermission.IMPORT,
+        (ast.Import, ast.ImportFrom),
+        '`import` is not allowed.',
+    )
+
+    super().generic_visit(node)
+
+
+def parse(
+    code: str,
+    permission: Optional[permissions.CodePermission] = None,
+) -> ast.AST:
+  try:
+    parsed_code = ast.parse(code, mode='exec')
+    if permission is not None:
+      _CodeValidator(code, permission).visit(parsed_code)
+  except SyntaxError as e:
+    raise errors.CodeError(code, e) from e
+  return parsed_code

pyglove/core/coding/parsing_test.py

@@ -0,0 +1,150 @@
+# 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 parsing
+from pyglove.core.coding import permissions
+
+
+class ParsePythonCodeTest(unittest.TestCase):
+
+  def test_no_permission_check(self):
+    ast = parsing.parse(inspect.cleandoc("""
+        if x > 0:
+          print(x)
+        """))
+    self.assertIsNotNone(ast)
+
+  def assert_allowed(self, code: str, permission: permissions.CodePermission):
+    ast = parsing.parse(inspect.cleandoc(code), permission)
+    self.assertIsNotNone(ast)
+
+  def assert_not_allowed(
+      self, code: str, permission: permissions.CodePermission
+  ):
+    with self.assertRaisesRegex(errors.CodeError, '.* is not allowed'):
+      parsing.parse(inspect.cleandoc(code), permission)
+
+  def test_parse_with_allowed_code(self):
+    self.assert_allowed(
+        """
+        x = 1
+        """,
+        permissions.CodePermission.ASSIGN,
+    )
+    self.assert_allowed(
+        """
+        if x > 0:
+          pass
+        """,
+        permissions.CodePermission.CONDITION,
+    )
+    self.assert_allowed(
+        """
+        for i in [1, 2, 3]:
+          pass
+        """,
+        permissions.CodePermission.LOOP,
+    )
+    self.assert_allowed(
+        """
+        foo(x=1, y=bar(2))
+        """,
+        permissions.CodePermission.CALL,
+    )
+    self.assert_allowed(
+        """
+        assert x > 1
+        """,
+        permissions.CodePermission.EXCEPTION,
+    )
+    self.assert_allowed(
+        """
+        class A:
+          pass
+        """,
+        permissions.CodePermission.CLASS_DEFINITION,
+    )
+    self.assert_allowed(
+        """
+        def foo(x, y):
+          return x + y
+        """,
+        permissions.CodePermission.FUNCTION_DEFINITION,
+    )
+    self.assert_allowed(
+        """
+        import re
+        """,
+        permissions.CodePermission.IMPORT,
+    )
+
+  def test_parse_with_not_allowed_code(self):
+    self.assert_not_allowed(
+        """
+        x = 1
+        """,
+        permissions.CodePermission.CONDITION,
+    )
+    self.assert_not_allowed(
+        """
+        if x > 0:
+          pass
+        """,
+        permissions.CodePermission.BASIC,
+    )
+    self.assert_not_allowed(
+        """
+        for i in range(5):
+          pass
+        """,
+        permissions.CodePermission.BASIC,
+    )
+    self.assert_not_allowed(
+        """
+        assert x > 1
+        """,
+        permissions.CodePermission.BASIC,
+    )
+    self.assert_not_allowed(
+        """
+        class A:
+          pass
+        """,
+        permissions.CodePermission.BASIC,
+    )
+    self.assert_not_allowed(
+        """
+        def foo(x, y):
+          return x + y
+        """,
+        permissions.CodePermission.BASIC,
+    )
+    self.assert_not_allowed(
+        """
+        import re
+        """,
+        permissions.CodePermission.BASIC,
+    )
+    self.assert_not_allowed(
+        """
+        range(5)
+        """,
+        permissions.CodePermission.ASSIGN,
+    )
+
+
+if __name__ == '__main__':
+  unittest.main()