langfun 0.1.2.dev202501150804__py3-none-any.whl → 0.1.2.dev202501160804__py3-none-any.whl

langfun/core/__init__.py

@@ -75,7 +75,6 @@ from langfun.core.sampling import random_sample
 from langfun.core.concurrent import RetryEntry
 from langfun.core.concurrent import concurrent_execute
 from langfun.core.concurrent import concurrent_map
-from langfun.core.concurrent import with_context_access
 from langfun.core.concurrent import with_retry
 
 # Interface for natural language formattable.

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``.

langfun/core/concurrent.py

@@ -25,7 +25,6 @@ import threading
 import time
 from typing import Annotated, Any, Callable, Iterable, Iterator, Literal, Sequence, Tuple, Type, Union
 
-from langfun.core import component
 import pyglove as pg
 
 
@@ -39,18 +38,6 @@ except ImportError:
   tqdm = None
 
 
-def with_context_access(func: Callable[..., Any]) -> Callable[..., Any]:
-  """Derives a user function with the access to the current context."""
-  with component.context() as current_context:
-    pass
-
-  def _func(*args, **kwargs) -> Any:
-    with component.context(**current_context):
-      return func(*args, **kwargs)
-
-  return _func
-
-
 class RetryError(RuntimeError):
   """Retry error."""
 
@@ -249,7 +236,8 @@ def concurrent_execute(
   try:
     executed_jobs = list(
         executor.map(
-            lambda job: job(), [with_context_access(job) for job in jobs]
+            lambda job: job(),
+            [pg.with_contextual_override(job) for job in jobs]
         )
     )
     for job in executed_jobs:
@@ -736,9 +724,7 @@ def concurrent_map(
         retry_interval=retry_interval,
         exponential_backoff=exponential_backoff,
     )
-    future = executor.submit(
-        with_context_access(job),
-    )
+    future = executor.submit(pg.with_contextual_override(job))
     pending_futures.append(future)
     future_to_job[future] = job
     total += 1

langfun/core/concurrent_test.py

@@ -29,22 +29,6 @@ class A(component.Component):
   y: int = component.contextual()
 
 
-class WithContextAccessTest(unittest.TestCase):
-
-  def test_context_access(self):
-    inputs = [A(1), A(2)]
-    with futures.ThreadPoolExecutor() as executor:
-      with component.context(y=3):
-        self.assertEqual(
-            list(
-                executor.map(
-                    concurrent.with_context_access(lambda x: x.y), inputs
-                )
-            ),
-            [3, 3],
-        )
-
-
 class RetryErrorTest(unittest.TestCase):
 
   def test_basics(self):

{langfun-0.1.2.dev202501150804.dist-info → langfun-0.1.2.dev202501160804.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: langfun
-Version: 0.1.2.dev202501150804
+Version: 0.1.2.dev202501160804
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.2.dev202501150804.dist-info → langfun-0.1.2.dev202501160804.dist-info}/RECORD

@@ -1,9 +1,9 @@
 langfun/__init__.py,sha256=fhfPXpHN7GoGqixpFfqhQkYxFs_siP_LhbjZhd3lhio,2497
-langfun/core/__init__.py,sha256=oo81OUA4Xy2q7icxcxAMicVzbla5gMDf9Nzw_iIAkis,4627
-langfun/core/component.py,sha256=HVrEoTL1Y01iqOHC3FYdbAOnffqfHHtGJXoK1vkdEwo,11583
-langfun/core/component_test.py,sha256=sG-T2wpvBfHqWGZE7sc4NayJj2aj5QFBzSwFiwrGEIc,10376
-langfun/core/concurrent.py,sha256=4CqOlH9YrtrK-Jijp09hLX7AYxRYL8slpy-tL2aqnIM,33121
-langfun/core/concurrent_test.py,sha256=0ucEYxPgTErOOK1TIN-TSqZphHS1Vl1VssyQP49cpOI,17997
+langfun/core/__init__.py,sha256=Cy3BU8R9VHnJmbxiN9QXXBb1K09ZoESYIJGf_uwJrDs,4571
+langfun/core/component.py,sha256=g1kQM0bryYYYWVDrSMnHfc74wIBbpfe5_B3s-UIP5GE,3028
+langfun/core/component_test.py,sha256=0CxTgjAud3aj8wBauFhG2FHDqrxCTl4OI4gzQTad-40,9254
+langfun/core/concurrent.py,sha256=zY-pXqlGqss_GI20tM1gXvyW8QepVPUuFNmutcIdhbI,32760
+langfun/core/concurrent_test.py,sha256=rc5T-2giWgtbwNuN6gmei7Uwo66HsJeeRtXZCpya_QU,17590
 langfun/core/console.py,sha256=V_mOiFi9oGh8gLsUeR56pdFDkuvYOpvQt7DY1KUTWTA,2535
 langfun/core/console_test.py,sha256=pBOcuNMJdVELywvroptfcRtJMsegMm3wSlHAL2TdxVk,1679
 langfun/core/langfunc.py,sha256=G50YgoVZ0y1GFw2ev41MlOqr6qa8YakbvNC0h_E0PiA,11140
@@ -146,8 +146,8 @@ langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fik
 langfun/core/templates/demonstration_test.py,sha256=SafcDQ0WgI7pw05EmPI2S4v1t3ABKzup8jReCljHeK4,2162
 langfun/core/templates/selfplay.py,sha256=yhgrJbiYwq47TgzThmHrDQTF4nDrTI09CWGhuQPNv-s,2273
 langfun/core/templates/selfplay_test.py,sha256=Ot__1P1M8oJfoTp-M9-PQ6HUXqZKyMwvZ5f7yQ3yfyM,2326
-langfun-0.1.2.dev202501150804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202501150804.dist-info/METADATA,sha256=HAqG3RU6S4kwv8Knx7vDVp01x8Kuc7It1QnYg0oeol8,8172
-langfun-0.1.2.dev202501150804.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-langfun-0.1.2.dev202501150804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202501150804.dist-info/RECORD,,
+langfun-0.1.2.dev202501160804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202501160804.dist-info/METADATA,sha256=_XM3ancZIb8-33gpRxLKmdJOBZsMfd1_2-4otzha19Q,8172
+langfun-0.1.2.dev202501160804.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+langfun-0.1.2.dev202501160804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202501160804.dist-info/RECORD,,