pyglove 0.4.5.dev202501140808__py3-none-any.whl → 0.4.5.dev202501160330__py3-none-any.whl

pyglove/core/__init__.py

@@ -125,6 +125,9 @@ compound_class = symbolic.compound_class
 # Method for declaring a boilerplated class from a symbolic instance.
 boilerplate_class = symbolic.boilerplate_class
 
+# Methods for contextual objects.
+ContextualObject = symbolic.ContextualObject
+contextual_attribute = symbolic.contextual_attribute
 
 #
 # Context manager for swapping wrapped class with their wrappers.
@@ -302,6 +305,10 @@ docstr = utils.docstr
 catch_errors = utils.catch_errors
 timeit = utils.timeit
 
+contextual_override = utils.contextual_override
+with_contextual_override = utils.with_contextual_override
+contextual_value = utils.contextual_value
+
 colored = utils.colored
 decolor = utils.decolor
 

pyglove/core/symbolic/__init__.py

@@ -86,6 +86,10 @@ from pyglove.core.symbolic.compounding import compound
 from pyglove.core.symbolic.compounding import compound_class
 from pyglove.core.symbolic.boilerplate import boilerplate_class
 
+from pyglove.core.symbolic.contextual_object import ContextualObject
+from pyglove.core.symbolic.contextual_object import ContextualAttribute
+from pyglove.core.symbolic.contextual_object import contextual_attribute
+
 # Inferential types.
 from pyglove.core.symbolic.base import Inferential
 

pyglove/core/symbolic/contextual_object.py

@@ -0,0 +1,288 @@
+# 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.
+"""Contextual objects.
+
+This module provides support for defining and working with `ContextualObject`s,
+whose attributes can dynamically adapt based on context managers and parent
+objects.
+
+# Overview:
+
+Contextual objects are specialized objects whose attribute values can be:
+1. Dynamically overridden using the `pg.contextual_override` context manager.
+2. Accessed via parent objects using the `pg.contextual_attribute` placeholder.
+
+This flexibility is particularly useful in scenarios where attributes need
+to respond to dynamic runtime conditions or inherit behavior from hierarchical
+structures.
+
+Example:
+
+```python
+class A(pg.ContextualObject):
+  x: int
+  y: Any = pg.contextual_attribute()
+
+a = A(1)
+print(a.x) # 1
+with pg.contextual_override(x=2):
+  print(a.x) # 2
+print(a.x) #1
+
+pg.Dict(y=3, a=a)
+print(a.y) # 3. (Accessing parent's y)
+```
+"""
+
+import threading
+from typing import Annotated, Any, ContextManager, Dict, Optional, Type
+from pyglove.core import utils as pg_utils
+from pyglove.core.symbolic import base
+from pyglove.core.symbolic import inferred as pg_inferred
+from pyglove.core.symbolic import object as pg_object
+from pyglove.core.views.html import tree_view
+
+
+class ContextualObject(pg_object.Object):
+  """Base class for contextual objects.
+
+  Contextual objects are objects whose attributes can be dynamically overridden
+  using `pg.contextual_override` or resolved through
+  `pg.contextual_attribute`, allowing them to inherit values from their
+  containing objects.
+
+  Usages:
+
+  ```
+  # Define a contextual object.
+  class A(pg.ContextualObject):
+    x: int
+    y: Any = pg.contextual_attribute()
+
+  # Create an instance of A
+  a = A(1)
+  print(a.x)  # Outputs: 1
+  print(a.y)  # Raises an error, as `a` has no containing object.
+
+  # Define another contextual object containing an instance of A
+  class B(pg.ContextualObject):
+    y: int
+    a: A
+
+  # Create an instance of B, containing "a"
+  b = B(y=2, a=a)
+  print(a.y)  # Outputs: 2, as "y" is resolved from the containing object (B).
+
+  # Contextual overrides are thread-specific
+  with pg.contextual_override(x=2):
+    print(a.x)  # Outputs: 2
+
+  # Thread-specific behavior of `pg.contextual_override`
+  def foo(a):
+    print(a.x)
+
+  with pg.contextual_override(x=3):
+    t = threading.Thread(target=foo, args=(a,))
+    t.start()
+    t.join()
+    # Outputs: 1, because `pg.contextual_override` is limited to the current
+    # thread to avoid clashes in multi-threaded environments.
+
+  # To propagate the override to a new thread, use `pg.with_contextual_override`
+  with pg.contextual_override(x=3):
+    t = threading.Thread(target=pg.with_contextual_override(foo), args=(a,))
+    t.start()
+    t.join()
+    # Outputs: 3, as the override is explicitly propagated.
+  ```
+  """
+
+  # Override __repr__ format to use inferred values when available.
+  __repr_format_kwargs__ = dict(
+      compact=True,
+      use_inferred=True,
+  )
+
+  # Override __str__ format to use inferred values when available.
+  __str_format_kwargs__ = dict(
+      compact=False,
+      verbose=False,
+      use_inferred=True,
+  )
+
+  def _on_bound(self):
+    super()._on_bound()
+    self._contextual_overrides = threading.local()
+
+  def _sym_inferred(self, key: str, **kwargs):
+    """Override to allow attribute to access scoped value.
+
+    Args:
+      key: attribute name.
+      **kwargs: Optional keyword arguments for value inference.
+
+    Returns:
+      The value of the symbolic attribute. If not available, returns the
+        default value.
+
+    Raises:
+      AttributeError: If the attribute does not exist or contextual attribute
+        is not ready.
+    """
+    if key not in self._sym_attributes:
+      raise AttributeError(key)
+
+    # Step 1: Try use value from `self.override`.
+    # The reason is that `self.override` is short-lived and explicitly specified
+    # by the user in scenarios like `LangFunc.render`, which should not be
+    # affected by `pg.contextual_override`.
+    v = pg_utils.contextual.get_scoped_value(self._contextual_overrides, key)
+    if v is not None:
+      return v.value
+
+    # Step 2: Try use value from `pg.contextual_override` with `override_attrs`.
+    # This gives users a chance to override the bound attributes of components
+    # from the top, allowing change of bindings without modifying the code
+    # that produces the components.
+    override = pg_utils.contextual.get_contextual_override(key)
+    if override and override.override_attrs:
+      return override.value
+
+    # Step 3: Try use value from the symbolic tree, starting from self to
+    # the root of the tree.
+    # Step 4: If the value is not present, use the value from `context()` (
+    # override_attrs=False).
+    # Step 5: Otherwise use the default value from `ContextualAttribute`.
+    return super()._sym_inferred(key, context_override=override, **kwargs)
+
+  def override(
+      self,
+      **kwargs
+  ) -> ContextManager[Dict[str, pg_utils.contextual.ContextualOverride]]:
+    """Context manager to override the attributes of this component."""
+    vs = {
+        k: pg_utils.contextual.ContextualOverride(v)
+        for k, v in kwargs.items()
+    }
+    return pg_utils.contextual.contextual_scope(
+        self._contextual_overrides, **vs
+    )
+
+  def __getattribute__(self, name: str) -> Any:
+    """Override __getattribute__ to deal with class attribute override."""
+    if not name.startswith('_') and hasattr(self.__class__, name):
+      tls = self.__dict__.get('_contextual_overrides', None)
+      if tls is not None:
+        v = pg_utils.contextual.get_scoped_value(tls, name)
+        if v is not None:
+          return v.value
+    return super().__getattribute__(name)
+
+
+class ContextualAttribute(
+    pg_inferred.ValueFromParentChain, tree_view.HtmlTreeView.Extension
+):
+  """Attributes whose values are inferred from the containing objects."""
+
+  NO_DEFAULT = (pg_utils.MISSING_VALUE,)
+
+  type: Annotated[Optional[Type[Any]], 'An optional type constraint.'] = None
+
+  default: Any = NO_DEFAULT
+
+  def value_from(
+      self,
+      parent,
+      *,
+      context_override: Optional[pg_utils.contextual.ContextualOverride] = None,
+      **kwargs,
+  ):
+    if (parent not in (None, self.sym_parent)
+        and isinstance(parent, ContextualObject)):
+      # Apply original search logic along the contextual object containing
+      # chain.
+      return super().value_from(parent, **kwargs)
+    elif parent is None:
+      # When there is no value inferred from the symbolic tree.
+      # Search context override, and then attribute-level default.
+      if context_override:
+        return context_override.value
+      if self.default == ContextualAttribute.NO_DEFAULT:
+        return pg_utils.MISSING_VALUE
+      return self.default
+    else:
+      return pg_utils.MISSING_VALUE
+
+  def _html_tree_view_content(
+      self,
+      *,
+      view: tree_view.HtmlTreeView,
+      parent: Any = None,
+      root_path: Optional[pg_utils.KeyPath] = None,
+      **kwargs,
+  ) -> tree_view.Html:
+    inferred_value = pg_utils.MISSING_VALUE
+    if isinstance(parent, base.Symbolic) and root_path:
+      inferred_value = parent.sym_inferred(
+          root_path.key, pg_utils.MISSING_VALUE
+      )
+
+    if inferred_value is not pg_utils.MISSING_VALUE:
+      kwargs.pop('name', None)
+      return view.render(
+          inferred_value, parent=self,
+          root_path=pg_utils.KeyPath('<inferred>', root_path),
+          **view.get_passthrough_kwargs(**kwargs)
+      )
+    return tree_view.Html.element(
+        'div',
+        [
+            '(not available)',
+        ],
+        css_classes=['unavailable-contextual'],
+    )
+
+  def _html_tree_view_config(self) -> Dict[str, Any]:
+    return tree_view.HtmlTreeView.get_kwargs(
+        super()._html_tree_view_config(),
+        dict(
+            collapse_level=1,
+        )
+    )
+
+  @classmethod
+  def _html_tree_view_css_styles(cls) -> list[str]:
+    return super()._html_tree_view_css_styles() + [
+        """
+        .contextual-attribute {
+          color: purple;
+        }
+        .unavailable-contextual {
+          color: gray;
+          font-style: italic;
+        }
+        """
+    ]
+
+
+# NOTE(daiyip): Returning Any instead of `pg.ContextualAttribute` to
+# avoid pytype check error as `contextual_attribute()` can be assigned to any
+# type.
+def contextual_attribute(
+    type: Optional[Type[Any]] = None,  # pylint: disable=redefined-builtin
+    default: Any = ContextualAttribute.NO_DEFAULT,
+) -> Any:
+  """Value marker for a contextual attribute."""
+  return ContextualAttribute(type=type, default=default, allow_partial=True)
+

pyglove/core/symbolic/contextual_object_test.py

@@ -0,0 +1,327 @@
+# 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
+from typing import Any
+import unittest
+import weakref
+
+from pyglove.core import utils as pg_utils
+from pyglove.core.symbolic import contextual_object
+
+contextual_override = pg_utils.contextual_override
+ContextualOverride = pg_utils.ContextualOverride
+get_contextual_override = pg_utils.get_contextual_override
+contextual_value = pg_utils.contextual_value
+all_contextual_values = pg_utils.all_contextual_values
+
+ContextualObject = contextual_object.ContextualObject
+contextual_attribute = contextual_object.contextual_attribute
+
+
+class ContextualObjectTest(unittest.TestCase):
+  """Tests for ContextualObject."""
+
+  def test_override(self):
+    class A(ContextualObject):
+      x: int
+
+    a = A(x=1)
+    with a.override(x=2, y=1):
+      self.assertEqual(a.x, 2)
+
+      # `y`` is not an attribute of `A`.
+      with self.assertRaises(AttributeError):
+        _ = a.y
+
+  def test_context(self):
+    class A(ContextualObject):
+      x: int
+      y: int = contextual_attribute()
+      z: int = contextual_attribute(default=-1)
+
+    with self.assertRaisesRegex(TypeError, '.* missing 1 required argument'):
+      _ = A()
+
+    a = A(x=1)
+    with self.assertRaisesRegex(
+        AttributeError, 'p'
+    ):
+      _ = a.p
+
+    with self.assertRaisesRegex(
+        AttributeError, '.* is not found under its context'
+    ):
+      _ = a.y
+
+    with contextual_override(y=1):
+      self.assertEqual(a.y, 1)
+
+    with self.assertRaisesRegex(
+        AttributeError, '.* is not found under its context'
+    ):
+      _ = a.y
+
+    # Use contextual default if it's not provided.
+    self.assertEqual(a.z, -1)
+
+    a1 = A(x=1, y=2)
+    self.assertEqual(a1.x, 1)
+    self.assertEqual(a1.y, 2)
+    self.assertEqual(a1.z, -1)
+
+    with contextual_override(x=3, y=3, z=3) as parent_override:
+      self.assertEqual(
+          parent_override,
+          dict(
+              x=ContextualOverride(3, cascade=False, override_attrs=False),
+              y=ContextualOverride(3, cascade=False, override_attrs=False),
+              z=ContextualOverride(3, cascade=False, override_attrs=False),
+          ),
+      )
+      self.assertEqual(
+          get_contextual_override('y'),
+          ContextualOverride(3, cascade=False, override_attrs=False),
+      )
+      self.assertEqual(contextual_value('x'), 3)
+      self.assertIsNone(contextual_value('f', None))
+      with self.assertRaisesRegex(KeyError, '.* does not exist'):
+        contextual_value('f')
+
+      self.assertEqual(all_contextual_values(), dict(x=3, y=3, z=3))
+
+      # Member attributes take precedence over `contextual_override`.
+      self.assertEqual(a1.x, 1)
+      self.assertEqual(a1.y, 2)
+
+      # Override attributes take precedence over member attribute.
+      with a1.override(y=3):
+        self.assertEqual(a1.y, 3)
+        with a1.override(y=4):
+          self.assertEqual(a1.y, 4)
+        self.assertEqual(a1.y, 3)
+      self.assertEqual(a1.y, 2)
+
+      # `contextual_override` takes precedence over contextual default.
+      self.assertEqual(a1.z, 3)
+
+      # Test nested contextual override with override_attrs=True (default).
+      with contextual_override(
+          y=4, z=4, override_attrs=True) as nested_override:
+        self.assertEqual(
+            nested_override,
+            dict(
+                x=ContextualOverride(3, cascade=False, override_attrs=False),
+                y=ContextualOverride(4, cascade=False, override_attrs=True),
+                z=ContextualOverride(4, cascade=False, override_attrs=True),
+            ),
+        )
+
+        # Member attribute is not overriden as current scope does not override
+        # `x``.
+        self.assertEqual(a1.x, 1)
+
+        # Member attribute is overriden.
+        self.assertEqual(a1.y, 4)
+
+        # `ContextualObject.override` takes precedence over
+        # `contextual_override(override_attrs=True)`.
+        with a1.override(y=3):
+          self.assertEqual(a1.y, 3)
+        self.assertEqual(a1.y, 4)
+
+        # Member default is overriden.
+        self.assertEqual(a1.z, 4)
+
+      self.assertEqual(a1.y, 2)
+      self.assertEqual(a1.z, 3)
+
+    self.assertEqual(a1.y, 2)
+    self.assertEqual(a1.z, -1)
+
+  def test_context_cascade(self):
+    class A(ContextualObject):
+      x: int
+      y: int = contextual_attribute()
+      z: int = contextual_attribute(default=-1)
+
+    a = A(1, 2)
+    self.assertEqual(a.x, 1)
+    self.assertEqual(a.y, 2)
+    self.assertEqual(a.z, -1)
+
+    with contextual_override(x=3, y=3, z=3, cascade=True):
+      self.assertEqual(a.x, 1)
+      self.assertEqual(a.y, 2)
+      self.assertEqual(a.z, 3)
+
+      # Outter `pg.contextual_override` takes precedence
+      # over inner `pg.contextual_override` when cascade=True.
+      with contextual_override(y=4, z=4, cascade=True):
+        self.assertEqual(a.x, 1)
+        self.assertEqual(a.y, 2)
+        self.assertEqual(a.z, 3)
+
+      with contextual_override(y=4, z=4, override_attrs=True):
+        self.assertEqual(a.x, 1)
+        self.assertEqual(a.y, 2)
+        self.assertEqual(a.z, 3)
+
+    self.assertEqual(a.x, 1)
+    self.assertEqual(a.y, 2)
+    self.assertEqual(a.z, -1)
+
+    with contextual_override(x=3, y=3, z=3, cascade=True, override_attrs=True):
+      self.assertEqual(a.x, 3)
+      self.assertEqual(a.y, 3)
+      self.assertEqual(a.z, 3)
+
+      with contextual_override(y=4, z=4, override_attrs=True):
+        self.assertEqual(a.x, 3)
+        self.assertEqual(a.y, 3)
+        self.assertEqual(a.z, 3)
+
+    self.assertEqual(a.x, 1)
+    self.assertEqual(a.y, 2)
+    self.assertEqual(a.z, -1)
+
+  def test_sym_inferred(self):
+    class A(ContextualObject):
+      x: int = 1
+      y: int = contextual_attribute()
+
+    a = A()
+    with self.assertRaisesRegex(
+        AttributeError, '.* is not found under its context'):
+      _ = a.sym_inferred('y')
+    self.assertIsNone(a.sym_inferred('y', default=None))
+
+    with self.assertRaises(AttributeError):
+      _ = a.sym_inferred('z')
+    self.assertIsNone(a.sym_inferred('z', default=None))
+
+  def test_weak_ref(self):
+    class A(ContextualObject):
+      x: int = 1
+
+    a = A()
+    self.assertIsNotNone(weakref.ref(a))
+
+
+class ContextualAttributeTest(unittest.TestCase):
+  """Tests for Component."""
+
+  def test_contextualibute_access(self):
+
+    class A(ContextualObject):
+      x: int
+      y: int = contextual_attribute()
+
+    # Not okay: `A.x` is required.
+    with self.assertRaisesRegex(TypeError, 'missing 1 required argument'):
+      _ = A()
+
+    # Okay: `A.y` is contextual.
+    a = A(1)
+
+    # `a.y` is not yet available from the context.
+    with self.assertRaises(AttributeError):
+      _ = a.y
+
+    class B(ContextualObject):
+      # Attributes with annotation will be treated as symbolic fields.
+      p: int
+      q: A = A(2)
+      z: int = contextual_attribute()
+
+    class C(ContextualObject):
+      a: int
+      b: B = B(2)
+      y: int = 1
+      z = 2
+
+    c = C(1)
+    b = c.b
+    a = b.q
+
+    # Test symbolic attributes declared from C.
+    self.assertTrue(c.sym_hasattr('a'))
+    self.assertTrue(c.sym_hasattr('b'))
+    self.assertTrue(c.sym_hasattr('y'))
+    self.assertFalse(c.sym_hasattr('z'))
+
+    # Contextual access to c.y from a.
+    self.assertEqual(a.y, 1)
+    self.assertEqual(b.z, 2)
+
+    # 'y' is not defined as an attribute in 'B'.
+    with self.assertRaises(AttributeError):
+      _ = b.y
+
+    c.rebind(y=2)
+    self.assertEqual(c.y, 2)
+    self.assertEqual(a.y, 2)
+
+    c.z = 3
+    self.assertEqual(c.z, 3)
+    self.assertEqual(b.z, 3)
+
+  def test_to_html(self):
+    class A(ContextualObject):
+      x: int = 1
+      y: int = contextual_attribute()
+
+    def assert_content(html, expected):
+      expected = inspect.cleandoc(expected).strip()
+      actual = html.content.strip()
+      if actual != expected:
+        print(actual)
+      self.assertEqual(actual.strip(), expected)
+
+    self.assertIn(
+        inspect.cleandoc(
+            """
+            .contextual-attribute {
+              color: purple;
+            }
+            .unavailable-contextual {
+              color: gray;
+              font-style: italic;
+            }
+            """
+        ),
+        A().to_html().style_section,
+    )
+
+    assert_content(
+        A().to_html(enable_summary_tooltip=False),
+        """
+        <details open class="pyglove a"><summary><div class="summary-title">A(...)</div></summary><div class="complex-value a"><details open class="pyglove int"><summary><div class="summary-name">x<span class="tooltip">x</span></div><div class="summary-title">int</div></summary><span class="simple-value int">1</span></details><details open class="pyglove contextual-attribute"><summary><div class="summary-name">y<span class="tooltip">y</span></div><div class="summary-title">ContextualAttribute(...)</div></summary><div class="unavailable-contextual">(not available)</div></details></div></details>
+        """
+    )
+
+    class B(ContextualObject):
+      z: Any
+      y: int = 2
+
+    b = B(A())
+    assert_content(
+        b.z.to_html(enable_summary_tooltip=False),
+        """
+        <details open class="pyglove a"><summary><div class="summary-title">A(...)</div></summary><div class="complex-value a"><details open class="pyglove int"><summary><div class="summary-name">x<span class="tooltip">x</span></div><div class="summary-title">int</div></summary><span class="simple-value int">1</span></details><details open class="pyglove contextual-attribute"><summary><div class="summary-name">y<span class="tooltip">y</span></div><div class="summary-title">ContextualAttribute(...)</div></summary><span class="simple-value int">2</span></details></div></details>
+        """
+    )
+
+
+if __name__ == '__main__':
+  unittest.main()

pyglove/core/utils/__init__.py

@@ -155,6 +155,14 @@ from pyglove.core.utils.error_utils import ErrorInfo
 from pyglove.core.utils.timing import timeit
 from pyglove.core.utils.timing import TimeIt
 
+# Value override from context manager.
+from pyglove.core.utils.contextual import ContextualOverride
+from pyglove.core.utils.contextual import contextual_override
+from pyglove.core.utils.contextual import with_contextual_override
+from pyglove.core.utils.contextual import get_contextual_override
+from pyglove.core.utils.contextual import contextual_value
+from pyglove.core.utils.contextual import all_contextual_values
+
 # Text color.
 from pyglove.core.utils.text_color import colored
 from pyglove.core.utils.text_color import colored_block

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/utils/timing.py

@@ -56,6 +56,16 @@ class TimeIt:
           **kwargs,
       )
 
+    def merge(self, other: 'TimeIt.Status') -> 'TimeIt.Status':
+      """Merges the status of two `pg.timeit`."""
+      assert other.name == self.name, (self.name, other.name)
+      return TimeIt.Status(
+          name=self.name,
+          elapse=self.elapse + other.elapse,
+          has_ended=self.has_ended and other.has_ended,
+          error=self.error or other.error,
+      )
+
   @dataclasses.dataclass
   class StatusSummary(json_conversion.JSONConvertible):
     """Aggregated summary for repeated calls for `pg.timeit`."""
@@ -125,7 +135,7 @@ class TimeIt:
     self._name: str = name
     self._start_time: Optional[float] = None
     self._end_time: Optional[float] = None
-    self._child_contexts: Dict[str, TimeIt] = {}
+    self._child_contexts: List[TimeIt] = []
     self._error: Optional[error_utils.ErrorInfo] = None
     self._parent: Optional[TimeIt] = None
 
@@ -137,13 +147,11 @@ class TimeIt:
   @property
   def children(self) -> List['TimeIt']:
     """Returns child contexts."""
-    return list(self._child_contexts.values())
+    return self._child_contexts
 
   def add(self, context: 'TimeIt'):
     """Adds a child context."""
-    if context.name in self._child_contexts:
-      raise ValueError(f'`timeit` with name {context.name!r} already exists.')
-    self._child_contexts[context.name] = context
+    self._child_contexts.append(context)
 
   def start(self):
     """Starts timing."""
@@ -206,11 +214,14 @@ class TimeIt:
             has_ended=self.has_ended, error=self._error,
         )
     }
-    for child in self._child_contexts.values():
+    for child in self._child_contexts:
       child_result = child.status()
       for k, v in child_result.items():
         key = f'{self.name}.{k}' if self.name else k
-        result[key] = v
+        if key in result:
+          result[key] = result[key].merge(v)
+        else:
+          result[key] = v
     return result
 
   def __enter__(self):

pyglove/core/utils/timing_test.py

@@ -62,9 +62,9 @@ class TimeItTest(unittest.TestCase):
         self.assertFalse(r['node.child'].has_ended)
         self.assertTrue(r['node.child.grandchild'].has_ended)
 
-        with self.assertRaisesRegex(ValueError, '.* already exists'):
-          with timing.timeit('grandchild'):
-            pass
+        with timing.timeit('grandchild') as t3:
+          time.sleep(0.5)
+        self.assertEqual(t1.children, [t2, t3])
 
     elapse2 = t.elapse
     self.assertTrue(t.has_ended)
@@ -72,17 +72,20 @@ class TimeItTest(unittest.TestCase):
     time.sleep(0.5)
     self.assertEqual(elapse2, t.elapse)
 
-    statuss = t.status()
+    status = t.status()
     self.assertEqual(
-        list(statuss.keys()),
+        list(status.keys()),
         ['node', 'node.child', 'node.child.grandchild']
     )
     self.assertEqual(
-        sorted([v.elapse for v in statuss.values()], reverse=True),
-        [v.elapse for v in statuss.values()],
+        status['node.child.grandchild'].elapse, t2.elapse + t3.elapse
+    )
+    self.assertEqual(
+        sorted([v.elapse for v in status.values()], reverse=True),
+        [v.elapse for v in status.values()],
     )
-    self.assertTrue(all(v.has_ended for v in statuss.values()))
-    self.assertFalse(any(v.has_error for v in statuss.values()))
+    self.assertTrue(all(v.has_ended for v in status.values()))
+    self.assertFalse(any(v.has_error for v in status.values()))
     status = t.status()
     json_dict = json_conversion.to_json(status)
     status2 = json_conversion.from_json(json_dict)

{pyglove-0.4.5.dev202501140808.dist-info → pyglove-0.4.5.dev202501160330.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pyglove
-Version: 0.4.5.dev202501140808
+Version: 0.4.5.dev202501160330
 Summary: PyGlove: A library for manipulating Python objects.
 Home-page: https://github.com/google/pyglove
 Author: PyGlove Authors

{pyglove-0.4.5.dev202501140808.dist-info → pyglove-0.4.5.dev202501160330.dist-info}/RECORD

@@ -1,5 +1,5 @@
 pyglove/__init__.py,sha256=YPUWALRDu4QuI7N3TUmVaVr513hUt0qXrjY582I1q5c,1352
-pyglove/core/__init__.py,sha256=629Ml4RzdBdtd_1Cz4HZoCde_ilOG2lAW02E90hStWw,9333
+pyglove/core/__init__.py,sha256=BF-8v1D5lTBAo4vlpDt0qyOG9NQ71oh_UfthBVrPaN4,9614
 pyglove/core/logging.py,sha256=Sr5nLyUQmc4CAEYAq8qbP3ghpjmpz2sOuhq2A0tgQ6I,2456
 pyglove/core/logging_test.py,sha256=3z_c6wnxbqDbwUmSOAZzeDPXvzoweYL5QHUQVMJ5Xgo,1917
 pyglove/core/coding/__init__.py,sha256=tuHIg19ZchtkOQbdFVTVLkUpBa5f1eo66XtnKw3lcIU,1645
@@ -65,7 +65,7 @@ pyglove/core/patching/pattern_based.py,sha256=UtSNB-ARNqVjXwZovjVi84QEoXUGLLBTgL
 pyglove/core/patching/pattern_based_test.py,sha256=PW1EcVfsFPB6wtgwg3s4dzvigWn3b5S8eMNGo0SJiZ0,2771
 pyglove/core/patching/rule_based.py,sha256=JAQp8mWeIOxwIdqusA3GmXia-fxQhQsxbUTmE329wF8,17038
 pyglove/core/patching/rule_based_test.py,sha256=qfy0ILmczV_LMHWEnwo2y079OrJsGYO0nKxSZdmIUcI,18782
-pyglove/core/symbolic/__init__.py,sha256=jYq-LwR1Ql3iMChjz9lN-0heDKiocmxR0ZJFEJ8RHHQ,5787
+pyglove/core/symbolic/__init__.py,sha256=JZvyaEcS1QxA8MaAGANBWMmRTQcGk_6F0kjFxjokJqg,6002
 pyglove/core/symbolic/base.py,sha256=aSZPSk7v3Niv6D1HfL_JEEXfI-U1cQnTxJxY9wXem54,77627
 pyglove/core/symbolic/base_test.py,sha256=yASIHIuWiUB1jf4nN-Y4XOjyvr8eQfRpr7s_1LZsu78,7188
 pyglove/core/symbolic/boilerplate.py,sha256=sQ3G25r5bo_UmIdjreL4jkAuQCXIHVlvUfGjjkNod6Y,5955
@@ -74,6 +74,8 @@ pyglove/core/symbolic/class_wrapper.py,sha256=xQiMh5vFlOQ76tbqsF5UWEghvU4d9UmvbN
 pyglove/core/symbolic/class_wrapper_test.py,sha256=GPgeLefIPj9hiD0ib6z2LyOm6A3O2f4-JHpE7dgfsNM,21644
 pyglove/core/symbolic/compounding.py,sha256=gvOodZ2gWHA0jNdwt8yvnRsPkHQDXDb5s88Uwp6DAqs,11703
 pyglove/core/symbolic/compounding_test.py,sha256=hOzrIROvajUTtPm0SUbEsEV4C1bXanhAoHinHrjZoXw,8320
+pyglove/core/symbolic/contextual_object.py,sha256=ar9Q_0P0HHbDf8kLHpS8GFZEMCRuCCB6DP18iGItiiw,9146
+pyglove/core/symbolic/contextual_object_test.py,sha256=wA5xfIhEHOC9qE3bbiA59CAPxWs9AVPaNiKEoprkWpQ,10209
 pyglove/core/symbolic/dict.py,sha256=7IxKSjUsGulUcNk3xqQnJ2ubGNUbqcc0ZFe8cKiigKA,36931
 pyglove/core/symbolic/dict_test.py,sha256=3MATsFw9Ci0jVPbAT2ivKVZ3PCnBBBvbkMWKgpgl-Fk,72915
 pyglove/core/symbolic/diff.py,sha256=zHfED0Bbq8G_HWNPj3vrOCWzt_062rFhx3BMlpCb9oo,16282
@@ -127,9 +129,11 @@ pyglove/core/typing/typed_missing.py,sha256=-l1omAu0jBZv5BnsFYXBqfvQwVBnmPh_X1wc
 pyglove/core/typing/typed_missing_test.py,sha256=TCNsb1SRpFaVdxYn2mB_yaLuja8w5Qn5NP7uGiZVBWs,2301
 pyglove/core/typing/value_specs.py,sha256=Yxdrz4aURMBrtqUW4to-2Vptc6Z6oqQrLyMBiAZt2bc,102302
 pyglove/core/typing/value_specs_test.py,sha256=MsScWDRaN_8pfRDO9MCp9HdUHVm_8wHWyKorWVSnhE4,127165
-pyglove/core/utils/__init__.py,sha256=cegrJtHy0h7DIUFcQN67TZJf9_f28OR-NrKbKKjPkeE,8183
+pyglove/core/utils/__init__.py,sha256=I2bRTzigU7qJVEATGlLUFkYzkiCBBCCEwrQyhsrRNmI,8602
 pyglove/core/utils/common_traits.py,sha256=PWxOgPhG5H60ZwfO8xNAEGRjFUqqDZQBWQYomOfvdy8,3640
 pyglove/core/utils/common_traits_test.py,sha256=DIuZB_1xfmeTVfWnGOguDQcDAM_iGgBOe8C-5CsIqBc,1122
+pyglove/core/utils/contextual.py,sha256=RxBQkDM2gB6QwZj_2oMels6oh-zQPGJJlinZbbqHuYQ,5148
+pyglove/core/utils/contextual_test.py,sha256=OOcthquVyAekTqt1RyYcEMHaockMIokpbv4pSf13Nzs,3252
 pyglove/core/utils/docstr_utils.py,sha256=5BY40kXozPKVGOB0eN8jy1P5_GHIzqFJ9FXAu_kzxaw,5119
 pyglove/core/utils/docstr_utils_test.py,sha256=i33VT6zHXEuIIJ4PPg1bVSfaYSnUsC8yST_NfpT-Uds,4228
 pyglove/core/utils/error_utils.py,sha256=ACIqtq4hsWrV00Gxd1OyeHCI-obGBR6J3H9VI1Y9clM,5780
@@ -146,8 +150,8 @@ pyglove/core/utils/text_color.py,sha256=xcCTCxY2qFNZs_jismMGus8scEXKBpYGAhpAgnz-
 pyglove/core/utils/text_color_test.py,sha256=3HtZpUB5XPr7A_5Fg4ZSMfNWeDRiQgSzmg9b1tctMI4,2801
 pyglove/core/utils/thread_local.py,sha256=i-CnyY3VREtLfAj4_JndBnsKuQLIgwG29ma8dAyRxbI,4839
 pyglove/core/utils/thread_local_test.py,sha256=AOqsdNsA-cYMvJckqxb25ql3Y5kDW0qLfBW1cl85Bnw,6757
-pyglove/core/utils/timing.py,sha256=fXdwEYXR84qZBC7hxbXbwztW2UKliJZLqCDvhz1qw_A,7241
-pyglove/core/utils/timing_test.py,sha256=OZxi551MJec9cKtpNbeep6jjzP68JUHhSmRHNKDFRMw,4928
+pyglove/core/utils/timing.py,sha256=gnHCA2IJ9DZ3Y0vb8ItB1Ap8ss8j_zN0OMUAdhoGy2Y,7550
+pyglove/core/utils/timing_test.py,sha256=t_awsYs5SqEO3_2u6HDbVxRcSxSfHsE4zMbFG-1ErZw,5011
 pyglove/core/utils/value_location.py,sha256=wAryIwQeEfrRddyGRk-KbLA7dnNDLhL-dSyFI9wIG5U,26756
 pyglove/core/utils/value_location_test.py,sha256=X6Gih3IoYugziwXWH8VGz5bPeb3Kq0CfZxwbNVIsZJo,21338
 pyglove/core/views/__init__.py,sha256=gll9ZBRYz4p_-LWOdzSR2a6UTWcJ8nR430trrP0yLCU,967
@@ -207,8 +211,8 @@ pyglove/ext/scalars/randoms.py,sha256=LkMIIx7lOq_lvJvVS3BrgWGuWl7Pi91-lA-O8x_gZs
 pyglove/ext/scalars/randoms_test.py,sha256=nEhiqarg8l_5EOucp59CYrpO2uKxS1pe0hmBdZUzRNM,2000
 pyglove/ext/scalars/step_wise.py,sha256=IDw3tuTpv0KVh7AN44W43zqm1-E0HWPUlytWOQC9w3Y,3789
 pyglove/ext/scalars/step_wise_test.py,sha256=TL1vJ19xVx2t5HKuyIzGoogF7N3Rm8YhLE6JF7i0iy8,2540
-pyglove-0.4.5.dev202501140808.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-pyglove-0.4.5.dev202501140808.dist-info/METADATA,sha256=_HkFXhfIYMiYBTTw47VFpp0lnLW5yojjf9NomI1d3II,7067
-pyglove-0.4.5.dev202501140808.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-pyglove-0.4.5.dev202501140808.dist-info/top_level.txt,sha256=wITzJSKcj8GZUkbq-MvUQnFadkiuAv_qv5qQMw0fIow,8
-pyglove-0.4.5.dev202501140808.dist-info/RECORD,,
+pyglove-0.4.5.dev202501160330.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+pyglove-0.4.5.dev202501160330.dist-info/METADATA,sha256=RpkNhAVoliHYNVVV3wCT6VzneEmNLys9KingM-qH6Cs,7067
+pyglove-0.4.5.dev202501160330.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+pyglove-0.4.5.dev202501160330.dist-info/top_level.txt,sha256=wITzJSKcj8GZUkbq-MvUQnFadkiuAv_qv5qQMw0fIow,8
+pyglove-0.4.5.dev202501160330.dist-info/RECORD,,