langfun 0.1.2.dev202501080804__py3-none-any.whl → 0.1.2.dev202501240804__py3-none-any.whl

langfun/core/component.py

@@ -13,10 +13,7 @@
 # limitations under the License.
 """langfun Component."""
 
-import contextlib
-import dataclasses
-import threading
-from typing import Annotated, Any, ContextManager, Iterator, Type
+from typing import ContextManager
 import pyglove as pg
 
 
@@ -24,22 +21,9 @@ import pyglove as pg
 RAISE_IF_HAS_ERROR = (pg.MISSING_VALUE,)
 
 
-class Component(pg.Object):
+class Component(pg.ContextualObject):
   """Base class for langfun components."""
 
-  # 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,
-  )
-
   # Allow symbolic assignment, which invalidates the object and recomputes
   # states upon update.
   allow_symbolic_assignment = True
@@ -75,122 +59,23 @@ class Component(pg.Object):
     if additional_fields:
       cls.update_schema(additional_fields)
 
-  def _on_bound(self):
-    super()._on_bound()
-    self._tls = 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 `lf.context`.
-    v = _get_scoped_value(self._tls, _CONTEXT_OVERRIDES, key)
-    if v is not None:
-      return v.value
-
-    # Step 2: Try use value from `lf.context` 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 = 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, 'ContextualOverride']]:
-    """Context manager to override the attributes of this component."""
-    vs = {k: ContextualOverride(v) for k, v in kwargs.items()}
-    return _contextual_scope(self._tls, _CONTEXT_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('_tls', None)
-      if tls is not None:
-        v = _get_scoped_value(tls, _CONTEXT_OVERRIDES, name)
-        if v is not None:
-          return v.value
-    return super().__getattribute__(name)
-
-
-_global_tls = threading.local()
-
-_CONTEXT_OVERRIDES = 'context_overrides'
-
 
-@dataclasses.dataclass(frozen=True)
-class ContextualOverride:
-  """Value marker for contextual override for an attribute."""
+# Aliases from PyGlove for ease of access.
+context = pg.contextual_override
+get_contextual_override = pg.utils.get_contextual_override
+context_value = pg.utils.contextual_value
+all_contextual_values = pg.utils.all_contextual_values
+contextual = pg.contextual_attribute
 
-  # Overridden value.
-  value: Any
-
-  # If True, this override will apply to both current scope and nested scope,
-  # meaning current `lf.context` will take precedence over all nested
-  # `lf.context` on this attribute.
-  cascade: bool = False
-
-  # If True, this override will apply to attributes that already have values.
-  override_attrs: bool = False
-
-
-def context(
-    *,
-    cascade: bool = False,
-    override_attrs: bool = False,
-    **variables,
-) -> ContextManager[dict[str, ContextualOverride]]:
-  """Context manager to provide overriden values for contextual attributes.
-
-  Args:
-    cascade: If True, this override will apply to both current scope and nested
-      scope, meaning that this `lf.context` will take precedence over all
-      nested `lf.context` 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_tls, _CONTEXT_OVERRIDES, **vs)
+# Decorator for setting the positional arguments for Component.
+use_init_args = pg.use_init_args
 
 
 def use_settings(
     *,
     cascade: bool = False,
     **settings,
-) -> ContextManager[dict[str, ContextualOverride]]:
+) -> ContextManager[dict[str, pg.utils.ContextualOverride]]:
   """Shortcut method for overriding component attributes.
 
   Args:
@@ -203,150 +88,3 @@ def use_settings(
     A dict of attribute names to their contextual overrides.
   """
   return context(cascade=cascade, override_attrs=True, **settings)
-
-
-def get_contextual_override(var_name: str) -> ContextualOverride | None:
-  """Returns the overriden contextual value in current scope."""
-  return _get_scoped_value(_global_tls, _CONTEXT_OVERRIDES, var_name)
-
-
-def context_value(var_name: str, default: Any = RAISE_IF_HAS_ERROR) -> Any:
-  """Returns the value of a variable defined in `lf.context`."""
-  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 contextual values provided from `lf.context` in scope."""
-  overrides = getattr(_global_tls, _CONTEXT_OVERRIDES, {})
-  return {k: v.value for k, v in overrides.items()}
-
-
-@contextlib.contextmanager
-def _contextual_scope(
-    tls: threading.local, tls_key, **variables
-) -> Iterator[dict[str, ContextualOverride]]:
-  """Context manager to set variables within a scope."""
-  previous_values = getattr(tls, tls_key, {})
-  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, current_values)
-    yield current_values
-  finally:
-    setattr(tls, tls_key, previous_values)
-
-
-def _get_scoped_value(
-    tls: threading.local, tls_key: str, var_name: str, default: Any = None
-) -> ContextualOverride:
-  """Gets the value for requested variable from current scope."""
-  scoped_values = getattr(tls, tls_key, {})
-  return scoped_values.get(var_name, default)
-
-
-class ContextualAttribute(
-    pg.symbolic.ValueFromParentChain, pg.views.HtmlTreeView.Extension
-):
-  """Attributes whose values are inferred from the context of the component.
-
-  Please see go/langfun-component#attribute-value-retrieval for details.
-  """
-
-  NO_DEFAULT = (pg.MISSING_VALUE,)
-
-  type: Annotated[Type[Any] | None, 'An optional type constraint.'] = None
-
-  default: Any = NO_DEFAULT
-
-  def value_from(
-      self,
-      parent,
-      *,
-      context_override: ContextualOverride | None = None,
-      **kwargs,
-  ):
-    if parent not in (None, self.sym_parent) and isinstance(parent, Component):
-      # Apply original search logic along the component 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.MISSING_VALUE
-      return self.default
-    else:
-      return pg.MISSING_VALUE
-
-  def _html_tree_view_content(
-      self,
-      *,
-      view: pg.views.HtmlTreeView,
-      parent: Any = None,
-      root_path: pg.KeyPath | None = None,
-      **kwargs,
-  ) -> pg.Html:
-    inferred_value = pg.MISSING_VALUE
-    if isinstance(parent, pg.Symbolic) and root_path:
-      inferred_value = parent.sym_inferred(root_path.key, pg.MISSING_VALUE)
-
-    if inferred_value is not pg.MISSING_VALUE:
-      kwargs.pop('name', None)
-      return view.render(
-          inferred_value, parent=self,
-          root_path=pg.KeyPath('<inferred>', root_path),
-          **view.get_passthrough_kwargs(**kwargs)
-      )
-    return pg.Html.element(
-        'div',
-        [
-            '(not available)',
-        ],
-        css_classes=['unavailable-contextual'],
-    )
-
-  def _html_tree_view_config(self) -> dict[str, Any]:
-    return pg.views.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 `lf.ContextualAttribute` to avoid
-# pytype check error as `contextual()` can be assigned to any type.
-def contextual(
-    type: Type[Any] | None = 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)
-
-
-# Decorator for setting the positional arguments for Component.
-use_init_args = pg.use_init_args

langfun/core/component_test.py

@@ -73,26 +73,7 @@ class ComponentContextTest(unittest.TestCase):
     self.assertEqual(a1.y, 2)
     self.assertEqual(a1.z, -1)
 
-    with lf.context(x=3, y=3, z=3) as parent_override:
-      self.assertEqual(
-          parent_override,
-          dict(
-              x=lf.ContextualOverride(3, cascade=False, override_attrs=False),
-              y=lf.ContextualOverride(3, cascade=False, override_attrs=False),
-              z=lf.ContextualOverride(3, cascade=False, override_attrs=False),
-          ),
-      )
-      self.assertEqual(
-          lf.get_contextual_override('y'),
-          lf.ContextualOverride(3, cascade=False, override_attrs=False),
-      )
-      self.assertEqual(lf.context_value('x'), 3)
-      self.assertIsNone(lf.context_value('f', None))
-      with self.assertRaisesRegex(KeyError, '.* does not exist'):
-        lf.context_value('f')
-
-      self.assertEqual(lf.all_contextual_values(), dict(x=3, y=3, z=3))
-
+    with lf.context(x=3, y=3, z=3):
       # Member attributes take precedence over `lf.context`.
       self.assertEqual(a1.x, 1)
       self.assertEqual(a1.y, 2)
@@ -109,15 +90,7 @@ class ComponentContextTest(unittest.TestCase):
       self.assertEqual(a1.z, 3)
 
       # Test nested contextual override with override_attrs=True (default).
-      with lf.context(y=4, z=4, override_attrs=True) as nested_override:
-        self.assertEqual(
-            nested_override,
-            dict(
-                x=lf.ContextualOverride(3, cascade=False, override_attrs=False),
-                y=lf.ContextualOverride(4, cascade=False, override_attrs=True),
-                z=lf.ContextualOverride(4, cascade=False, override_attrs=True),
-            ),
-        )
+      with lf.context(y=4, z=4, override_attrs=True):
 
         # Member attribute is not overriden as current scope does not override
         # `x``.