pyglove 0.4.5.dev202411132359__py3-none-any.whl → 0.4.5.dev202501250807__py3-none-any.whl

pyglove/core/utils/contextual.py

@@ -0,0 +1,147 @@
+# 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.
+"""Injecting and manipulating values through context managers."""
+
+import contextlib
+import dataclasses
+import threading
+from typing import Any, Callable, ContextManager, Iterator, Optional
+
+from pyglove.core.utils import missing
+
+RAISE_IF_HAS_ERROR = (missing.MISSING_VALUE,)
+_TLS_KEY_CONTEXTUAL_OVERRIDES = '__contextual_overrides__'
+_global_contextual_overrides = threading.local()
+
+
+@dataclasses.dataclass(frozen=True)
+class ContextualOverride:
+  """Value marker for contextual override for an attribute."""
+
+  # Overridden value.
+  value: Any
+
+  # If True, this override will apply to both current scope and nested scope,
+  # meaning current `pg.contextual_override` will take precedence over all
+  # nested `pg.contextual_override` on this attribute.
+  cascade: bool = False
+
+  # If True, this override will apply to attributes that already have values.
+  override_attrs: bool = False
+
+
+def contextual_override(
+    *,
+    cascade: bool = False,
+    override_attrs: bool = False,
+    **variables,
+) -> ContextManager[dict[str, ContextualOverride]]:
+  """Context manager to provide contextual values under a scope.
+
+  Please be aware that contextual value override are per-thread. If you want
+  to propagate the contextual value override to other threads, please obtain
+  a wrapper function for a user function using
+  `pg.with_contextual_override(func)`.
+
+  Args:
+    cascade: If True, this override will apply to both current scope and nested
+      scope, meaning that this `pg.contextual_override` will take precedence
+      over all nested `pg.contextual_override` on the overriden variables.
+    override_attrs: If True, this override will apply to attributes that already
+      have values. Otherwise overridden variables will only be used for
+      contextual attributes whose values are not present.
+    **variables: Key/values as override for contextual attributes.
+
+  Returns:
+    A dict of attribute names to their contextual overrides.
+  """
+  vs = {}
+  for k, v in variables.items():
+    if not isinstance(v, ContextualOverride):
+      v = ContextualOverride(v, cascade, override_attrs)
+    vs[k] = v
+  return contextual_scope(_global_contextual_overrides, **vs)
+
+
+def with_contextual_override(func: Callable[..., Any]) -> Callable[..., Any]:
+  """Wraps a user function with the access to the current contextual override.
+
+  The wrapped function can be called from another thread.
+
+  Args:
+    func: The user function to be wrapped.
+
+  Returns:
+    A wrapper function that have the access to the current contextual override,
+    which can be called from another thread.
+  """
+  with contextual_override() as current_context:
+    pass
+
+  def _func(*args, **kwargs) -> Any:
+    with contextual_override(**current_context):
+      return func(*args, **kwargs)
+
+  return _func
+
+
+def get_contextual_override(var_name: str) -> Optional[ContextualOverride]:
+  """Returns the overriden contextual value in current scope."""
+  return get_scoped_value(_global_contextual_overrides, var_name)
+
+
+def contextual_value(var_name: str, default: Any = RAISE_IF_HAS_ERROR) -> Any:
+  """Returns the value of a variable defined in `pg.contextual_override`."""
+  override = get_contextual_override(var_name)
+  if override is None:
+    if default == RAISE_IF_HAS_ERROR:
+      raise KeyError(f'{var_name!r} does not exist in current context.')
+    return default
+  return override.value
+
+
+def all_contextual_values() -> dict[str, Any]:
+  """Returns all values provided from `pg.contextual_override` in scope."""
+  overrides = getattr(
+      _global_contextual_overrides, _TLS_KEY_CONTEXTUAL_OVERRIDES, {}
+  )
+  return {k: v.value for k, v in overrides.items()}
+
+
+@contextlib.contextmanager
+def contextual_scope(
+    tls: threading.local, **variables
+) -> Iterator[dict[str, ContextualOverride]]:
+  """Context manager to set variables within a scope."""
+  previous_values = getattr(tls, _TLS_KEY_CONTEXTUAL_OVERRIDES, {})
+  current_values = dict(previous_values)
+  for k, v in variables.items():
+    old_v = current_values.get(k, None)
+    if old_v and old_v.cascade:
+      v = old_v
+    current_values[k] = v
+  try:
+    setattr(tls, _TLS_KEY_CONTEXTUAL_OVERRIDES, current_values)
+    yield current_values
+  finally:
+    setattr(tls, _TLS_KEY_CONTEXTUAL_OVERRIDES, previous_values)
+
+
+def get_scoped_value(
+    tls: threading.local, var_name: str, default: Any = None
+) -> ContextualOverride:
+  """Gets the value for requested variable from current scope."""
+  scoped_values = getattr(tls, _TLS_KEY_CONTEXTUAL_OVERRIDES, {})
+  return scoped_values.get(var_name, default)
+

pyglove/core/utils/contextual_test.py

@@ -0,0 +1,88 @@
+# 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 concurrent.futures
+import unittest
+from pyglove.core.utils import contextual
+
+
+class ContextualTest(unittest.TestCase):
+
+  def test_contextual_override(self):
+    with contextual.contextual_override(x=3, y=3, z=3) as parent_override:
+      self.assertEqual(
+          parent_override,
+          dict(
+              x=contextual.ContextualOverride(
+                  3, cascade=False, override_attrs=False
+              ),
+              y=contextual.ContextualOverride(
+                  3, cascade=False, override_attrs=False
+              ),
+              z=contextual.ContextualOverride(
+                  3, cascade=False, override_attrs=False
+              ),
+          ),
+      )
+      self.assertEqual(
+          contextual.get_contextual_override('y'),
+          contextual.ContextualOverride(3, cascade=False, override_attrs=False),
+      )
+      self.assertEqual(contextual.contextual_value('x'), 3)
+      self.assertIsNone(contextual.contextual_value('f', None))
+      with self.assertRaisesRegex(KeyError, '.* does not exist'):
+        contextual.contextual_value('f')
+
+      self.assertEqual(contextual.all_contextual_values(), dict(x=3, y=3, z=3))
+
+      # Test nested contextual override with override_attrs=True (default).
+      with contextual.contextual_override(
+          y=4, z=4, override_attrs=True) as nested_override:
+        self.assertEqual(
+            nested_override,
+            dict(
+                x=contextual.ContextualOverride(
+                    3, cascade=False, override_attrs=False
+                ),
+                y=contextual.ContextualOverride(
+                    4, cascade=False, override_attrs=True
+                ),
+                z=contextual.ContextualOverride(
+                    4, cascade=False, override_attrs=True
+                ),
+            ),
+        )
+
+    # Test nested contextual override with cascade=True.
+    with contextual.contextual_override(x=3, y=3, z=3, cascade=True):
+      with contextual.contextual_override(y=4, z=4, cascade=True):
+        self.assertEqual(contextual.contextual_value('x'), 3)
+        self.assertEqual(contextual.contextual_value('y'), 3)
+        self.assertEqual(contextual.contextual_value('z'), 3)
+
+  def test_with_contextual_override(self):
+    def func(i):
+      del i
+      return contextual.contextual_value('x')
+
+    pool = concurrent.futures.ThreadPoolExecutor()
+    with contextual.contextual_override(x=3):
+      self.assertEqual(contextual.with_contextual_override(func)(0), 3)
+      self.assertEqual(
+          list(pool.map(contextual.with_contextual_override(func), range(1))),
+          [3]
+      )
+
+
+if __name__ == '__main__':
+  unittest.main()

pyglove/core/{object_utils → utils}/docstr_utils_test.py

@@ -11,11 +11,9 @@
 # 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.docstr_utils."""
-
 import inspect
 import unittest
-from pyglove.core.object_utils import docstr_utils
+from pyglove.core.utils import docstr_utils
 
 
 class DocStrTest(unittest.TestCase):

pyglove/core/{object_utils → utils}/error_utils.py

@@ -21,8 +21,8 @@ import sys
 import traceback
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type, Union
 
-from pyglove.core.object_utils import formatting
-from pyglove.core.object_utils import json_conversion
+from pyglove.core.utils import formatting
+from pyglove.core.utils import json_conversion
 
 
 @dataclasses.dataclass(frozen=True)
@@ -102,7 +102,7 @@ def catch_errors(
 
   Examples::
 
-    with pg.object_utils.catch_errors(
+    with pg.utils.catch_errors(
         [
             RuntimeErrror,
             (ValueError, 'Input is wrong.')

pyglove/core/{object_utils → utils}/error_utils_test.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import inspect
 import unittest
-from pyglove.core.object_utils import error_utils
+from pyglove.core.utils import error_utils
 
 
 class ErrorInfoTest(unittest.TestCase):

pyglove/core/{object_utils → utils}/formatting.py

@@ -18,7 +18,7 @@ import enum
 import io
 import sys
 from typing import Any, Callable, ContextManager, Dict, List, Optional, Sequence, Set, Tuple
-from pyglove.core.object_utils import thread_local
+from pyglove.core.utils import thread_local
 
 
 _TLS_STR_FORMAT_KWARGS = '_str_format_kwargs'

pyglove/core/{object_utils → utils}/formatting_test.py

@@ -11,11 +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.object_utils.formatting."""
 import inspect
 import unittest
 
-from pyglove.core.object_utils import formatting
+from pyglove.core.utils import formatting
 
 
 class Foo(formatting.Formattable):

pyglove/core/{object_utils → utils}/hierarchical.py

@@ -14,9 +14,9 @@
 """Operating hierarchical object."""
 
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
-from pyglove.core.object_utils import common_traits
-from pyglove.core.object_utils.missing import MISSING_VALUE
-from pyglove.core.object_utils.value_location import KeyPath
+from pyglove.core.utils import common_traits
+from pyglove.core.utils.missing import MISSING_VALUE
+from pyglove.core.utils.value_location import KeyPath
 
 
 def traverse(value: Any,
@@ -33,7 +33,7 @@ def traverse(value: Any,
       print(path)
 
     tree = {'a': [{'c': [1, 2]}, {'d': {'g': (3, 4)}}], 'b': 'foo'}
-    pg.object_utils.traverse(tree, preorder_visit)
+    pg.utils.traverse(tree, preorder_visit)
 
     # Should print:
     # 'a'
@@ -48,10 +48,10 @@ def traverse(value: Any,
 
   Args:
     value: A maybe hierarchical value to traverse.
-    preorder_visitor_fn: Preorder visitor function.
-      Function signature is (path, value) -> should_continue.
-    postorder_visitor_fn: Postorder visitor function.
-      Function signature is (path, value) -> should_continue.
+    preorder_visitor_fn: Preorder visitor function. Function signature is (path,
+      value) -> should_continue.
+    postorder_visitor_fn: Postorder visitor function. Function signature is
+      (path, value) -> should_continue.
     root_path: The key path of the root value.
 
   Returns:
@@ -111,7 +111,7 @@ def transform(value: Any,
         'e': 'bar',
         'f': 4
     }
-    output = pg.object_utils.transform(inputs, _remove_int)
+    output = pg.utils.transform(inputs, _remove_int)
     assert output == {
         'a': {
             'c': ['bar'],
@@ -123,11 +123,11 @@ def transform(value: Any,
   Args:
     value: Any python value type. If value is a list of dict, transformation
       will occur recursively.
-    transform_fn: Transform function in signature
-      (path, value) -> new value
-      If new value is MISSING_VALUE, key will be deleted.
+    transform_fn: Transform function in signature (path, value) -> new value If
+      new value is MISSING_VALUE, key will be deleted.
     root_path: KeyPath of the root.
     inplace: If True, perform transformation in place.
+
   Returns:
     Transformed value.
   """
@@ -186,7 +186,7 @@ def flatten(src: Any, flatten_complex_keys: bool = True) -> Any:
         'b': 'hi',
         'c': None
     }
-    output = pg.object_utils.flatten(inputs)
+    output = pg.utils.flatten(inputs)
     assert output == {
         'a.e': 1,
         'a.f[0].g': 2,
@@ -200,9 +200,9 @@ def flatten(src: Any, flatten_complex_keys: bool = True) -> Any:
   Args:
     src: source value to flatten.
     flatten_complex_keys: if True, complex keys such as 'x.y' will be flattened
-    as 'x'.'y'. For example:
-      {'a': {'b.c': 1}} will be flattened into {'a.b.c': 1} if this flag is on,
-      otherwise it will be flattened as {'a[b.c]': 1}.
+      as 'x'.'y'. For example: {'a': {'b.c': 1}} will be flattened into
+      {'a.b.c': 1} if this flag is on, otherwise it will be flattened as
+      {'a[b.c]': 1}.
 
   Returns:
     For primitive value types, `src` itself will be returned.
@@ -464,7 +464,7 @@ def merge(value_list: List[Any],
             'f': 10
         }
     }
-    output = pg.object_utils.merge([original, patch])
+    output = pg.utils.merge([original, patch])
     assert output == {
         'a': 1,
         # b is updated.
@@ -486,14 +486,12 @@ def merge(value_list: List[Any],
       value. The merge process will keep input values intact.
     merge_fn: A function to handle value merge that will be called for updated
       or added keys. If a branch is added/updated, the root of branch will be
-      passed to merge_fn.
-      the signature of function is:
-      `(path, left_value, right_value) -> final_value`
-      If a key is only present in src dict, old_value is MISSING_VALUE;
-      If a key is only present in dest dict, new_value is MISSING_VALUE;
-      otherwise both new_value and old_value are filled.
-      If final_value is MISSING_VALUE for a path, it will be removed from its
-      parent collection.
+      passed to merge_fn. the signature of function is: `(path, left_value,
+      right_value) -> final_value` If a key is only present in src dict,
+      old_value is MISSING_VALUE; If a key is only present in dest dict,
+      new_value is MISSING_VALUE; otherwise both new_value and old_value are
+      filled. If final_value is MISSING_VALUE for a path, it will be removed
+      from its parent collection.
 
   Returns:
     A merged value.

pyglove/core/{object_utils → utils}/hierarchical_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.object_utils.hierarchical."""
-
 import unittest
-from pyglove.core.object_utils import common_traits
-from pyglove.core.object_utils import hierarchical
-from pyglove.core.object_utils import value_location
+from pyglove.core.utils import common_traits
+from pyglove.core.utils import hierarchical
+from pyglove.core.utils import value_location
 
 
 class TraverseTest(unittest.TestCase):

pyglove/core/{object_utils → utils}/json_conversion.py

@@ -37,7 +37,7 @@ JSONPrimitiveType = Union[int, float, bool, str]
 # pytype doesn't support recursion. Use Any instead of 'JSONValueType'
 # in List and Dict.
 JSONListType = List[Any]
-JSONDictType = Dict[str, Any]
+JSONDictType = Dict[Union[str, int], Any]
 JSONValueType = Union[JSONPrimitiveType, JSONListType, JSONDictType]
 
 # pylint: enable=invalid-name

pyglove/core/{object_utils → utils}/json_conversion_test.py

@@ -11,13 +11,11 @@
 # 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.json_conversion."""
-
 import abc
 import typing
 import unittest
-from pyglove.core.object_utils import json_conversion
 from pyglove.core.typing import inspect as pg_inspect
+from pyglove.core.utils import json_conversion
 
 
 class X:

pyglove/core/{object_utils → utils}/missing.py

@@ -14,8 +14,8 @@
 """Representing missing value for a field."""
 
 from typing import Any, Dict
-from pyglove.core.object_utils import formatting
-from pyglove.core.object_utils import json_conversion
+from pyglove.core.utils import formatting
+from pyglove.core.utils import json_conversion
 
 
 class MissingValue(formatting.Formattable, json_conversion.JSONConvertible):

pyglove/core/{object_utils → utils}/missing_test.py

@@ -11,11 +11,9 @@
 # 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.missing."""
-
 import unittest
-from pyglove.core.object_utils import json_conversion
-from pyglove.core.object_utils import missing
+from pyglove.core.utils import json_conversion
+from pyglove.core.utils import missing
 
 
 class MissingValueTest(unittest.TestCase):

pyglove/core/utils/text_color.py

@@ -0,0 +1,128 @@
+# 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.
+"""Utility library for text coloring."""
+
+import re
+from typing import List, Optional
+
+try:
+  import termcolor      # pylint: disable=g-import-not-at-top
+except ImportError:
+  termcolor = None
+
+
+# Regular expression for ANSI color characters.
+_ANSI_COLOR_REGEX = re.compile(r'\x1b\[[0-9;]*m')
+
+
+def decolor(text: str) -> str:
+  """De-colors a string that may contains ANSI color characters."""
+  return re.sub(_ANSI_COLOR_REGEX, '', text)
+
+
+def colored(
+    text: str,
+    color: Optional[str] = None,
+    background: Optional[str] = None,
+    styles: Optional[List[str]] = None,
+) -> str:
+  """Returns the colored text with ANSI color characters.
+
+  Args:
+    text: A string that may or may not already has ANSI color characters.
+    color: A string for text colors. Applicable values are:
+      'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'.
+    background: A string for background colors. Applicable values are:
+      'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'.
+    styles: A list of strings for applying styles on the text.
+      Applicable values are:
+      'bold', 'dark', 'underline', 'blink', 'reverse', 'concealed'.
+
+  Returns:
+    A string with ANSI color characters embracing the entire text.
+  """
+  if not termcolor:
+    return text
+  return termcolor.colored(
+      text,
+      color=color,
+      on_color=('on_' + background) if background else None,
+      attrs=styles
+  )
+
+
+def colored_block(
+    text: str,
+    block_start: str,
+    block_end: str,
+    color: Optional[str] = None,
+    background: Optional[str] = None,
+    styles: Optional[List[str]] = None,
+) -> str:
+  """Apply colors to text blocks.
+
+  Args:
+    text: A string that may or may not already has ANSI color characters.
+    block_start: A string that signals the start of a block. E.g. '{{'
+    block_end: A string that signals the end of a block. E.g. '}}'.
+    color: A string for text colors. Applicable values are:
+      'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'.
+    background: A string for background colors. Applicable values are:
+      'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'.
+    styles: A list of strings for applying styles on the text.
+      Applicable values are:
+      'bold', 'dark', 'underline', 'blink', 'reverse', 'concealed'.
+
+  Returns:
+    A string with ANSI color characters embracing the matched text blocks.
+  """
+  if not color and not background and not styles:
+    return text
+
+  s = []
+  start_index = 0
+  end_index = 0
+  previous_color = None
+
+  def write_nonblock_text(text: str, previous_color: Optional[str]):
+    if previous_color:
+      s.append(previous_color)
+    s.append(text)
+
+  while start_index < len(text):
+    start_index = text.find(block_start, end_index)
+    if start_index == -1:
+      write_nonblock_text(text[end_index:], previous_color)
+      break
+
+    # Deal with text since last block.
+    since_last_block = text[end_index:start_index]
+    write_nonblock_text(since_last_block, previous_color)
+    colors = re.findall(_ANSI_COLOR_REGEX, since_last_block)
+    if colors:
+      previous_color = colors[-1]
+
+    # Match block.
+    end_index = text.find(block_end, start_index + len(block_start))
+    if end_index == -1:
+      write_nonblock_text(text[start_index:], previous_color)
+      break
+    end_index += len(block_end)
+
+    # Write block text.
+    block = text[start_index:end_index]
+    block = colored(
+        block, color=color, background=background, styles=styles)
+    s.append(block)
+  return ''.join(s)