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

pyglove/core/symbolic/base_test.py

@@ -11,15 +11,13 @@
 # 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.symbolic.base."""
-
 import copy
 import inspect
 from typing import Any
 import unittest
 
-from pyglove.core import object_utils
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core import views
 from pyglove.core.symbolic import base
 from pyglove.core.symbolic.dict import Dict
@@ -33,7 +31,7 @@ class FieldUpdateTest(unittest.TestCase):
   def test_basics(self):
     x = Dict(x=1)
     f = pg_typing.Field('x', pg_typing.Int())
-    update = base.FieldUpdate(object_utils.KeyPath('x'), x, f, 1, 2)
+    update = base.FieldUpdate(utils.KeyPath('x'), x, f, 1, 2)
     self.assertEqual(update.path, 'x')
     self.assertIs(update.target, x)
     self.assertIs(update.field, f)
@@ -42,15 +40,15 @@ class FieldUpdateTest(unittest.TestCase):
 
   def test_format(self):
     self.assertEqual(
-        base.FieldUpdate(
-            object_utils.KeyPath('x'), Dict(x=1), None, 1, 2
-        ).format(compact=True),
+        base.FieldUpdate(utils.KeyPath('x'), Dict(x=1), None, 1, 2).format(
+            compact=True
+        ),
         'FieldUpdate(parent_path=, path=x, old_value=1, new_value=2)',
     )
 
     self.assertEqual(
         base.FieldUpdate(
-            object_utils.KeyPath('a'), Dict(x=Dict(a=1)).x, None, 1, 2
+            utils.KeyPath('a'), Dict(x=Dict(a=1)).x, None, 1, 2
         ).format(compact=True),
         'FieldUpdate(parent_path=x, path=a, old_value=1, new_value=2)',
     )
@@ -59,34 +57,34 @@ class FieldUpdateTest(unittest.TestCase):
     x = Dict()
     f = pg_typing.Field('x', pg_typing.Int())
     self.assertEqual(
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 2),
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 2),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 2),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 2),
     )
 
     # Targets are not the same instance.
     self.assertNotEqual(
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 2),
-        base.FieldUpdate(object_utils.KeyPath('a'), Dict(), f, 1, 2),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 2),
+        base.FieldUpdate(utils.KeyPath('a'), Dict(), f, 1, 2),
     )
 
     # Fields are not the same instance.
     self.assertNotEqual(
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 2),
-        base.FieldUpdate(object_utils.KeyPath('b'), x, copy.copy(f), 1, 2),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 2),
+        base.FieldUpdate(utils.KeyPath('b'), x, copy.copy(f), 1, 2),
     )
 
     self.assertNotEqual(
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 2),
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 0, 2),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 2),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 0, 2),
     )
 
     self.assertNotEqual(
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 2),
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 1),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 2),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 1),
     )
 
     self.assertNotEqual(
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 2), Dict()
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 2), Dict()
     )
 
 

pyglove/core/symbolic/boilerplate.py

@@ -15,11 +15,10 @@
 
 import copy
 import inspect
-
 from typing import Any, List, Optional, Type
 
-from pyglove.core import object_utils
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.symbolic import flags
 from pyglove.core.symbolic import object as pg_object
 
@@ -129,9 +128,9 @@ def boilerplate_class(
   cls.auto_register = True
 
   allow_partial = value.allow_partial
-  def _freeze_field(path: object_utils.KeyPath,
-                    field: pg_typing.Field,
-                    value: Any) -> Any:
+  def _freeze_field(
+      path: utils.KeyPath, field: pg_typing.Field, value: Any
+  ) -> Any:
     # We do not do validation since Object is already in valid form.
     del path
     if not isinstance(field.key, pg_typing.ListKey):

pyglove/core/symbolic/class_wrapper.py

@@ -26,8 +26,8 @@ import types
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type, Union
 
 from pyglove.core import detouring
-from pyglove.core import object_utils
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 
 from pyglove.core.symbolic import dict as pg_dict  # pylint: disable=unused-import
 from pyglove.core.symbolic import list as pg_list  # pylint: disable=unused-import
@@ -45,11 +45,7 @@ class ClassWrapperMeta(pg_object.ObjectMeta):
 
   def __getattr__(self, name):
     """Pass through attribute requests to sym_wrapped_cls."""
-    try:
-      return super().__getattr__(name)
-    except AttributeError:
-      wrapped_cls = object.__getattribute__(self, 'sym_wrapped_cls')
-      return getattr(wrapped_cls, name)
+    return getattr(object.__getattribute__(self, 'sym_wrapped_cls'), name)
 
 
 class ClassWrapper(pg_object.Object, metaclass=ClassWrapperMeta):
@@ -75,7 +71,7 @@ class _SubclassedWrapperBase(ClassWrapper):
   # the `__init__` method.
   auto_typing = False
 
-  @object_utils.explicit_method_override
+  @utils.explicit_method_override
   def __init__(self, *args, **kwargs):
     """Overridden __init__ to construct symbolic wrapper only."""
     # NOTE(daiyip): We avoid `__init__` to be called multiple times.
@@ -104,7 +100,7 @@ class _SubclassedWrapperBase(ClassWrapper):
   def __init_subclass__(cls):
     # Class wrappers inherit `__init__` from the user class. Therefore, we mark
     # all of them as explicitly overridden.
-    object_utils.explicit_method_override(cls.__init__)
+    utils.explicit_method_override(cls.__init__)
 
     super().__init_subclass__()
     if cls.__init__ is _SubclassedWrapperBase.__init__:
@@ -133,7 +129,7 @@ class _SubclassedWrapperBase(ClassWrapper):
       init_arg_list, arg_fields = _extract_init_signature(
           cls, auto_doc=cls.auto_doc, auto_typing=cls.auto_typing)
 
-      @object_utils.explicit_method_override
+      @utils.explicit_method_override
       @functools.wraps(cls.__init__)
       def _sym_init(self, *args, **kwargs):
         _SubclassedWrapperBase.__init__(self, *args, **kwargs)
@@ -526,7 +522,7 @@ def apply_wrappers(
   """
   if not wrapper_classes:
     wrapper_classes = []
-    for _, c in object_utils.JSONConvertible.registered_types():
+    for _, c in utils.JSONConvertible.registered_types():
       if (issubclass(c, ClassWrapper)
           and c not in (ClassWrapper, _SubclassedWrapperBase)
           and (not where or where(c))
@@ -548,13 +544,13 @@ def _extract_init_signature(
     # Read args docstr from both class doc string and __init__ doc string.
     args_docstr = dict()
     if cls.__doc__:
-      cls_docstr = object_utils.DocStr.parse(cls.__doc__)
+      cls_docstr = utils.DocStr.parse(cls.__doc__)
       args_docstr = cls_docstr.args
     if init_method.__doc__:
-      init_docstr = object_utils.DocStr.parse(init_method.__doc__)
+      init_docstr = utils.DocStr.parse(init_method.__doc__)
       args_docstr.update(init_docstr.args)
-    docstr = object_utils.DocStr(
-        object_utils.DocStrStyle.GOOGLE,
+    docstr = utils.DocStr(
+        utils.DocStrStyle.GOOGLE,
         short_description=None,
         long_description=None,
         examples=[],

pyglove/core/symbolic/class_wrapper_test.py

@@ -588,7 +588,7 @@ class ClassWrapperTest(unittest.TestCase):
     self.assertIsInstance(C(1, 2), ClassWrapper)
     self.assertTrue(pg_eq(C(1, 2), C(1, 2)))
     self.assertEqual(list(C.__schema__.fields.keys()), ['x', 'y'])
-    self.assertEqual(repr(C), f'<class {C.type_name!r}>')
+    self.assertEqual(repr(C), f'<class {C.__type_name__!r}>')
 
   def test_custom_metaclass(self):
 
@@ -604,7 +604,7 @@ class ClassWrapperTest(unittest.TestCase):
     A1 = pg_wrap(A)  # pylint: disable=invalid-name
     self.assertTrue(issubclass(A1, ClassWrapper))
     self.assertTrue(issubclass(A1, A))
-    self.assertEqual(A1.type_name, 'pyglove.core.symbolic.class_wrapper_test.A')
+    self.assertEqual(A1.__type_name__, 'pyglove.core.symbolic.class_wrapper_test.A')
     self.assertEqual(A1.__schema__, pg_typing.Schema([]))
     self.assertEqual(A1.foo, 'foo')
     self.assertRegex(repr(A1), r'Symbolic\[.*\]')

pyglove/core/symbolic/compounding.py

@@ -19,7 +19,7 @@ import sys
 import types
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
-from pyglove.core import object_utils
+from pyglove.core import utils
 from pyglove.core.symbolic.base import Symbolic
 from pyglove.core.symbolic.object import Object
 import pyglove.core.typing as pg_typing
@@ -39,7 +39,7 @@ class Compound(Object):
     # from the user class to compound with.
     Object.__init_subclass__(cls)
 
-  @object_utils.explicit_method_override
+  @utils.explicit_method_override
   def __init__(self, *args, **kwargs):
     # `explicit_init` allows the `__init__` of the other classes that sit after
     # `Compound` to be bypassed.

pyglove/core/symbolic/compounding_test.py

@@ -11,15 +11,13 @@
 # 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.compounding."""
-
 import abc
 import dataclasses
 import sys
 import unittest
 
-from pyglove.core import object_utils
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.symbolic.compounding import compound as pg_compound
 from pyglove.core.symbolic.compounding import compound_class as pg_compound_class
 from pyglove.core.symbolic.dict import Dict
@@ -145,7 +143,7 @@ class UserClassTest(unittest.TestCase):
     class A(Object):
       x: int
 
-      @object_utils.explicit_method_override
+      @utils.explicit_method_override
       def __init__(self, x):
         super().__init__(x=x)
         assert type(self) is A  # pylint: disable=unidiomatic-typecheck

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)
+