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

pyglove/core/coding/permissions.py

@@ -0,0 +1,100 @@
+# 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 permissions."""
+
+import contextlib
+import enum
+from typing import Optional
+
+from pyglove.core import utils
+
+
+class CodePermission(enum.Flag):
+  """Permissions for code execution."""
+
+  # Allows assignment.
+  ASSIGN = enum.auto()
+
+  # Allows conditions.
+  CONDITION = enum.auto()
+
+  # Allows loops.
+  LOOP = enum.auto()
+
+  # Call functions or methods.
+  CALL = enum.auto()
+
+  # Allows exception.
+  EXCEPTION = enum.auto()
+
+  # Allows class definitions.
+  CLASS_DEFINITION = enum.auto()
+
+  # Allows function definitions.
+  FUNCTION_DEFINITION = enum.auto()
+
+  # Allows import.
+  IMPORT = enum.auto()
+
+  @classmethod
+  @property
+  def BASIC(cls) -> 'CodePermission':    # pylint: disable=invalid-name
+    """Returns basic permissions."""
+    return CodePermission.ASSIGN | CodePermission.CALL
+
+  @classmethod
+  @property
+  def ALL(cls) -> 'CodePermission':    # pylint: disable=invalid-name
+    """Returns all permissions."""
+    return (
+        CodePermission.BASIC | CodePermission.CONDITION | CodePermission.LOOP |
+        CodePermission.EXCEPTION | CodePermission.CLASS_DEFINITION |
+        CodePermission.FUNCTION_DEFINITION | CodePermission.IMPORT)
+
+
+_TLS_CODE_RUN_PERMISSION = '__code_run_permission__'
+
+
+@contextlib.contextmanager
+def permission(perm: CodePermission):
+  """Context manager for controling the permission for code execution.
+
+  When the `permission` context manager is nested, the outtermost permission
+  will be used. This design allows users to control permission at the top level.
+
+  Args:
+    perm: Code execution permission.
+
+  Yields:
+    Actual permission applied.
+  """
+
+  outter_perm = utils.thread_local_get(_TLS_CODE_RUN_PERMISSION, None)
+
+  # Use the top-level permission as the actual permission
+  if outter_perm is not None:
+    perm = outter_perm
+
+  utils.thread_local_set(_TLS_CODE_RUN_PERMISSION, perm)
+
+  try:
+    yield perm
+  finally:
+    if outter_perm is None:
+      utils.thread_local_del(_TLS_CODE_RUN_PERMISSION)
+
+
+def get_permission() -> Optional[CodePermission]:
+  """Gets the current permission for code execution."""
+  return utils.thread_local_get(_TLS_CODE_RUN_PERMISSION, None)

pyglove/core/coding/permissions_test.py

@@ -0,0 +1,93 @@
+# 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 unittest
+from pyglove.core.coding import permissions
+
+
+class CodePermissionTest(unittest.TestCase):
+
+  def assert_set(
+      self,
+      permission: permissions.CodePermission,
+      flag: permissions.CodePermission,
+  ):
+    self.assertEqual(permission & flag, flag)
+
+  def assert_not_set(
+      self,
+      permission: permissions.CodePermission,
+      flag: permissions.CodePermission,
+  ):
+    self.assertFalse(permission & flag)
+
+  def test_basic(self):
+    self.assert_set(
+        permissions.CodePermission.BASIC, permissions.CodePermission.ASSIGN
+    )
+    self.assert_set(
+        permissions.CodePermission.BASIC, permissions.CodePermission.CALL
+    )
+    self.assert_set(
+        permissions.CodePermission.BASIC, permissions.CodePermission.CALL
+    )
+
+  def test_all(self):
+    self.assert_set(
+        permissions.CodePermission.ALL, permissions.CodePermission.BASIC
+    )
+    self.assert_set(
+        permissions.CodePermission.ALL, permissions.CodePermission.CONDITION
+    )
+    self.assert_set(
+        permissions.CodePermission.ALL, permissions.CodePermission.LOOP
+    )
+    self.assert_set(
+        permissions.CodePermission.ALL, permissions.CodePermission.EXCEPTION
+    )
+    self.assert_set(
+        permissions.CodePermission.ALL,
+        permissions.CodePermission.CLASS_DEFINITION,
+    )
+    self.assert_set(
+        permissions.CodePermission.ALL,
+        permissions.CodePermission.FUNCTION_DEFINITION,
+    )
+    self.assert_set(
+        permissions.CodePermission.ALL, permissions.CodePermission.IMPORT
+    )
+
+  def test_xor(self):
+    self.assert_not_set(
+        permissions.CodePermission.ALL ^ permissions.CodePermission.BASIC,
+        permissions.CodePermission.BASIC,
+    )
+    self.assert_set(
+        permissions.CodePermission.ALL ^ permissions.CodePermission.BASIC,
+        permissions.CodePermission.CONDITION,
+    )
+
+  def test_permission_control(self):
+    self.assertIsNone(permissions.get_permission())
+    with permissions.permission(permissions.CodePermission.BASIC):
+      self.assertEqual(
+          permissions.get_permission(), permissions.CodePermission.BASIC
+      )
+      with permissions.permission(permissions.CodePermission.ALL):
+        self.assertEqual(
+            permissions.get_permission(), permissions.CodePermission.BASIC
+        )
+
+
+if __name__ == '__main__':
+  unittest.main()

pyglove/core/geno/base.py

@@ -19,9 +19,9 @@ import random
 import types
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
-from pyglove.core import object_utils
 from pyglove.core import symbolic
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 
 
 class AttributeDict(dict):
@@ -37,12 +37,20 @@ class AttributeDict(dict):
 
 
 @symbolic.members([
-    ('location',
-     pg_typing.Object(object_utils.KeyPath, default=object_utils.KeyPath()),
-     ('KeyPath of associated genetic encoder relative to parent object '
-      'template. This allows DNA generator to apply rule based on locations.')),
-    ('hints',
-     pg_typing.Any(default=None), 'Hints for DNA generator to consume.')
+    (
+        'location',
+        pg_typing.Object(utils.KeyPath, default=utils.KeyPath()),
+        (
+            'KeyPath of associated genetic encoder relative to parent object'
+            ' template. This allows DNA generator to apply rule based on'
+            ' locations.'
+        ),
+    ),
+    (
+        'hints',
+        pg_typing.Any(default=None),
+        'Hints for DNA generator to consume.',
+    ),
 ])
 class DNASpec(symbolic.Object):
   """Base class for DNA specifications (genotypes).
@@ -175,7 +183,7 @@ class DNASpec(symbolic.Object):
     """Returns all decision points in their declaration order."""
 
   @property
-  def decision_ids(self) -> List[object_utils.KeyPath]:
+  def decision_ids(self) -> List[utils.KeyPath]:
     """Returns decision IDs."""
     return list(self._decision_point_by_id.keys())
 
@@ -286,7 +294,7 @@ class DNASpec(symbolic.Object):
     return self.parent_spec if self.is_space else self.parent_spec.parent_choice
 
   @property
-  def id(self) -> object_utils.KeyPath:
+  def id(self) -> utils.KeyPath:
     """Returns a path of locations from the root as the ID for current node."""
     if self._id is None:
       parent = self.parent_spec
@@ -295,18 +303,20 @@ class DNASpec(symbolic.Object):
       elif self.is_space:
         assert parent.is_categorical, parent
         assert self.index is not None
-        self._id = object_utils.KeyPath(
-            ConditionalKey(self.index, len(parent.candidates)),
-            parent.id) + self.location
+        self._id = (
+            utils.KeyPath(
+                ConditionalKey(self.index, len(parent.candidates)), parent.id
+            )
+            + self.location
+        )
       else:
         # Float() or a multi-choice spec of a parent Choice.
         self._id = parent.id + self.location
     return self._id
 
-  def get(self,
-          name_or_id: Union[object_utils.KeyPath, str],
-          default: Any = None
-          ) -> Union['DecisionPoint', List['DecisionPoint']]:
+  def get(
+      self, name_or_id: Union[utils.KeyPath, str], default: Any = None
+  ) -> Union['DecisionPoint', List['DecisionPoint']]:
     """Get decision point(s) by name or ID."""
     try:
       return self[name_or_id]
@@ -314,9 +324,8 @@ class DNASpec(symbolic.Object):
       return default
 
   def __getitem__(
-      self,
-      name_or_id: Union[object_utils.KeyPath, str]
-      ) -> Union['DecisionPoint', List['DecisionPoint']]:
+      self, name_or_id: Union[utils.KeyPath, str]
+  ) -> Union['DecisionPoint', List['DecisionPoint']]:
     """Get decision point(s) by name or ID ."""
     v = self._named_decision_points.get(name_or_id, None)
     if v is None:
@@ -475,7 +484,7 @@ class DNA(symbolic.Object):
   # Allow assignment on symbolic attributes.
   allow_symbolic_assignment = True
 
-  @object_utils.explicit_method_override
+  @utils.explicit_method_override
   def __init__(
       self,
       value: Union[None, int, float, str, List[Any], Tuple[Any]] = None,
@@ -727,7 +736,7 @@ class DNA(symbolic.Object):
     return self._decision_by_id_cache
 
   @property
-  def decision_ids(self) -> List[object_utils.KeyPath]:
+  def decision_ids(self) -> List[utils.KeyPath]:
     """Returns decision IDs."""
     self._ensure_dna_spec()
     return self._spec.decision_ids
@@ -1249,9 +1258,11 @@ class DNA(symbolic.Object):
     return dna
 
   def to_numbers(
-      self, flatten: bool = True,
-      ) -> Union[List[Union[int, float, str]],
-                 object_utils.Nestable[Union[int, float, str]]]:
+      self,
+      flatten: bool = True,
+  ) -> Union[
+      List[Union[int, float, str]], utils.Nestable[Union[int, float, str]]
+  ]:
     """Returns a (maybe) nested structure of numbers as decisions.
 
     Args:
@@ -1338,7 +1349,7 @@ class DNA(symbolic.Object):
               f'Location: {dna_spec.location.path}.')
         children = []
         for i, choice in enumerate(decision):
-          choice_location = object_utils.KeyPath(i, dna_spec.location)
+          choice_location = utils.KeyPath(i, dna_spec.location)
           if not isinstance(choice, int):
             raise ValueError(
                 f'Choice value should be int. Encountered: {choice}, '
@@ -1410,7 +1421,7 @@ class DNA(symbolic.Object):
 
     if type_info:
       json_value = {
-          object_utils.JSONConvertible.TYPE_NAME_KEY: (
+          utils.JSONConvertible.TYPE_NAME_KEY: (
               self.__class__.__serialization_key__
           ),
           'format': 'compact',
@@ -1435,7 +1446,8 @@ class DNA(symbolic.Object):
       json_value: Dict[str, Any],
       *,
       allow_partial: bool = False,
-      root_path: Optional[object_utils.KeyPath] = None) -> 'DNA':
+      root_path: Optional[utils.KeyPath] = None,
+  ) -> 'DNA':
     """Class method that load a DNA from a JSON value.
 
     Args:
@@ -1472,8 +1484,8 @@ class DNA(symbolic.Object):
     return not self.children
 
   def __getitem__(
-      self, key: Union[int, slice, str, object_utils.KeyPath, 'DecisionPoint']
-      ) -> Union[None, 'DNA', List[Optional['DNA']]]:
+      self, key: Union[int, slice, str, utils.KeyPath, 'DecisionPoint']
+  ) -> Union[None, 'DNA', List[Optional['DNA']]]:
     """Get an immediate child DNA or DNA in the sub-tree.
 
     Args:
@@ -1504,10 +1516,11 @@ class DNA(symbolic.Object):
         v = self._decision_by_id[key]
       return v
 
-  def get(self,
-          key: Union[int, slice, str, object_utils.KeyPath, 'DecisionPoint'],
-          default: Any = None
-          ) -> Union[Any, None, 'DNA', List[Optional['DNA']]]:
+  def get(
+      self,
+      key: Union[int, slice, str, utils.KeyPath, 'DecisionPoint'],
+      default: Any = None,
+  ) -> Union[Any, None, 'DNA', List[Optional['DNA']]]:
     """Get an immediate child DNA or DNA in the sub-tree."""
     try:
       return self[key]
@@ -1529,8 +1542,9 @@ class DNA(symbolic.Object):
           return True
       else:
         raise ValueError(
-            f'DNA.__contains__ does not accept '
-            f'{object_utils.quote_if_str(dna_or_value)!r}.')
+            'DNA.__contains__ does not accept '
+            f'{utils.quote_if_str(dna_or_value)!r}.'
+        )
     return False
 
   def __hash__(self):
@@ -1599,21 +1613,20 @@ class DNA(symbolic.Object):
       verbose: bool = True,
       root_indent: int = 0,
       *,
-      markdown: bool = False,
       list_wrap_threshold: int = 80,
       as_dict: bool = False,
       **kwargs,
   ):
     """Customize format method for DNA for more compact representation."""
     if as_dict and self.spec:
-      details = object_utils.format(
+      details = utils.format(
           self.to_dict(value_type='choice_and_literal'),
           False,
           verbose,
           root_indent,
-          **kwargs)
+          **kwargs,
+      )
       s = f'DNA({details})'
-      compact = False
     else:
       if 'list_wrap_threshold' not in kwargs:
         kwargs['list_wrap_threshold'] = list_wrap_threshold
@@ -1623,7 +1636,7 @@ class DNA(symbolic.Object):
       elif self.is_leaf:
         s = f'DNA({self.value!r})'
       else:
-        rep = object_utils.format(
+        rep = utils.format(
             self.to_json(compact=True, type_info=False),
             compact,
             verbose,
@@ -1636,7 +1649,7 @@ class DNA(symbolic.Object):
           s = f'DNA{rep}'
         else:
           s = f'DNA({rep})'
-    return object_utils.maybe_markdown_quote(s, markdown)
+    return s
 
   def parameters(
       self, use_literal_values: bool = False) -> Dict[str, str]:

pyglove/core/geno/base_test.py

@@ -11,12 +11,10 @@
 # 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.geno.DNA."""
-
 import unittest
 
-from pyglove.core import object_utils
 from pyglove.core import symbolic
+from pyglove.core import utils
 from pyglove.core.geno.base import ConditionalKey
 from pyglove.core.geno.base import DNA
 from pyglove.core.geno.categorical import manyof
@@ -1144,7 +1142,7 @@ class ConditionalKeyTest(unittest.TestCase):
     self.assertEqual(key.num_choices, 5)
 
   def test_to_str(self):
-    key = object_utils.KeyPath(['a', ConditionalKey(1, 5), 'b'])
+    key = utils.KeyPath(['a', ConditionalKey(1, 5), 'b'])
     self.assertEqual(str(key), 'a[=1/5].b')
 
 

pyglove/core/geno/categorical.py

@@ -18,10 +18,9 @@ import re
 import types
 from typing import Any, List, Optional, Union
 
-from pyglove.core import object_utils
 from pyglove.core import symbolic
 from pyglove.core import typing as pg_typing
-
+from pyglove.core import utils
 from pyglove.core.geno.base import DecisionPoint
 from pyglove.core.geno.base import DNA
 from pyglove.core.geno.base import DNASpec
@@ -139,14 +138,15 @@ class Choices(DecisionPoint):
       for i in range(self.num_choices):
         subchoice_spec = Choices(
             subchoice_index=i,
-            location=object_utils.KeyPath(i),
+            location=utils.KeyPath(i),
             num_choices=1,
             candidates=self.candidates,
             literal_values=self.literal_values,
             distinct=self.distinct,
             sorted=self.sorted,
             name=self.name,
-            hints=self.hints)
+            hints=self.hints,
+        )
         self._decision_points.extend(subchoice_spec.decision_points)
         subchoice_specs.append(subchoice_spec)
       self._subchoice_specs = symbolic.List(subchoice_specs)
@@ -158,7 +158,8 @@ class Choices(DecisionPoint):
         self._decision_points.extend(c.decision_points)
 
   def _update_children_paths(
-      self, old_path: object_utils.KeyPath, new_path: object_utils.KeyPath):
+      self, old_path: utils.KeyPath, new_path: utils.KeyPath
+  ):
     """Trigger path change for subchoices so their IDs can be invalidated."""
     super()._update_children_paths(old_path, new_path)
     if self._subchoice_specs:
@@ -337,7 +338,7 @@ class Choices(DecisionPoint):
               f'DNA child values should be sorted. '
               f'Encountered: {sub_dna_values}, Location: {self.location.path}.')
       for i, sub_dna in enumerate(dna):
-        sub_location = object_utils.KeyPath(i, self.location)
+        sub_location = utils.KeyPath(i, self.location)
         if not isinstance(sub_dna.value, int):
           raise ValueError(
               f'Choice value should be int. Encountered: {sub_dna.value}, '
@@ -598,16 +599,21 @@ class Choices(DecisionPoint):
         s.append(_indent(f'({i}): {value_str}\n', root_indent + 1))
     s.append(_indent(']', root_indent))
     if show_id:
-      kvlist = [('id', object_utils.quote_if_str(str(self.id)), '\'\'')]
+      kvlist = [('id', str(self.id), '\'\'')]
     else:
       kvlist = []
-    additionl_properties = object_utils.kvlist_str(kvlist + [
-        ('name', object_utils.quote_if_str(self.name), None),
-        ('distinct', self.distinct, True),
-        ('sorted', self.sorted, False),
-        ('hints', object_utils.quote_if_str(self.hints), None),
-        ('subchoice_index', self.subchoice_index, None)
-    ], compact=False, root_indent=root_indent)
+    additionl_properties = utils.kvlist_str(
+        kvlist
+        + [
+            ('name', self.name, None),
+            ('distinct', self.distinct, True),
+            ('sorted', self.sorted, False),
+            ('hints', self.hints, None),
+            ('subchoice_index', self.subchoice_index, None),
+        ],
+        compact=False,
+        root_indent=root_indent,
+    )
     if additionl_properties:
       s.append(', ')
       s.append(additionl_properties)
@@ -615,14 +621,16 @@ class Choices(DecisionPoint):
     return ''.join(s)
 
 
-def manyof(num_choices: int,
-           candidates: List[DNASpec],
-           distinct: bool = True,
-           sorted: bool = False,      # pylint: disable=redefined-builtin
-           literal_values: Optional[List[Union[str, int, float]]] = None,
-           hints: Any = None,
-           location: Union[str, object_utils.KeyPath] = object_utils.KeyPath(),
-           name: Optional[str] = None) -> Choices:
+def manyof(
+    num_choices: int,
+    candidates: List[DNASpec],
+    distinct: bool = True,
+    sorted: bool = False,  # pylint: disable=redefined-builtin
+    literal_values: Optional[List[Union[str, int, float]]] = None,
+    hints: Any = None,
+    location: Union[str, utils.KeyPath] = utils.KeyPath(),
+    name: Optional[str] = None,
+) -> Choices:
   """Returns a multi-choice specification.
 
   It creates the genotype for :func:`pyglove.manyof`.
@@ -674,11 +682,13 @@ def manyof(num_choices: int,
                  hints=hints, location=location, name=name)
 
 
-def oneof(candidates: List[DNASpec],
-          literal_values: Optional[List[Union[str, int, float]]] = None,
-          hints: Any = None,
-          location: Union[str, object_utils.KeyPath] = object_utils.KeyPath(),
-          name: Optional[str] = None) -> Choices:
+def oneof(
+    candidates: List[DNASpec],
+    literal_values: Optional[List[Union[str, int, float]]] = None,
+    hints: Any = None,
+    location: Union[str, utils.KeyPath] = utils.KeyPath(),
+    name: Optional[str] = None,
+) -> Choices:
   """Returns a single choice specification.
 
   It creates the genotype for :func:`pyglove.oneof`.
@@ -716,4 +726,3 @@ def oneof(candidates: List[DNASpec],
   """
   return manyof(1, candidates, literal_values=literal_values,
                 hints=hints, location=location, name=name)
-

pyglove/core/geno/custom.py

@@ -17,10 +17,9 @@ import random
 import types
 from typing import Any, Callable, List, Optional, Union
 
-from pyglove.core import object_utils
 from pyglove.core import symbolic
 from pyglove.core import typing as pg_typing
-
+from pyglove.core import utils
 from pyglove.core.geno.base import DecisionPoint
 from pyglove.core.geno.base import DNA
 
@@ -125,7 +124,7 @@ class CustomDecisionPoint(DecisionPoint):
           f'CustomDecisionPoint expects string type DNA. '
           f'Encountered: {dna!r}, Location: {self.location.path}.')
 
-  def sym_jsonify(self, **kwargs: Any) -> object_utils.JSONValueType:
+  def sym_jsonify(self, **kwargs: Any) -> utils.JSONValueType:
     """Overrides sym_jsonify to exclude non-serializable fields."""
     exclude_keys = kwargs.pop('exclude_keys', [])
     exclude_keys.extend(['random_dna_fn', 'next_dna_fn'])
@@ -150,24 +149,28 @@ class CustomDecisionPoint(DecisionPoint):
     if not compact:
       return super().format(compact, verbose, root_indent, **kwargs)
     if show_id:
-      kvlist = [('id', object_utils.quote_if_str(str(self.id)), '\'\'')]
+      kvlist = [('id', str(self.id), '\'\'')]
     else:
       kvlist = []
-    details = object_utils.kvlist_str(kvlist + [
-        ('hyper_type', object_utils.quote_if_str(self.hyper_type), None),
-        ('name', object_utils.quote_if_str(self.name), None),
-        ('hints', object_utils.quote_if_str(self.hints), None),
-    ])
+    details = utils.kvlist_str(
+        kvlist
+        + [
+            ('hyper_type', self.hyper_type, None),
+            ('name', self.name, None),
+            ('hints', self.hints, None),
+        ]
+    )
     return f'{self.__class__.__name__}({details})'
 
 
-def custom(hyper_type: Optional[str] = None,
-           next_dna_fn: Optional[
-               Callable[[Optional[DNA]], Optional[DNA]]] = None,
-           random_dna_fn: Optional[Callable[[Any], DNA]] = None,
-           hints: Any = None,
-           location: object_utils.KeyPath = object_utils.KeyPath(),
-           name: Optional[str] = None) -> CustomDecisionPoint:
+def custom(
+    hyper_type: Optional[str] = None,
+    next_dna_fn: Optional[Callable[[Optional[DNA]], Optional[DNA]]] = None,
+    random_dna_fn: Optional[Callable[[Any], DNA]] = None,
+    hints: Any = None,
+    location: utils.KeyPath = utils.KeyPath(),
+    name: Optional[str] = None,
+) -> CustomDecisionPoint:
   """Returns a custom decision point.
 
   It creates the genotype for subclasses of :func:`pyglove.hyper.CustomHyper`.