pyglove 0.4.5.dev202501150808__tar.gz → 0.4.5.dev202501160330__tar.gz

{pyglove-0.4.5.dev202501150808 → pyglove-0.4.5.dev202501160330}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pyglove
-Version: 0.4.5.dev202501150808
+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.dev202501150808 → pyglove-0.4.5.dev202501160330}/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-0.4.5.dev202501150808 → pyglove-0.4.5.dev202501160330}/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-0.4.5.dev202501160330/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-0.4.5.dev202501160330/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-0.4.5.dev202501150808 → pyglove-0.4.5.dev202501160330}/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