Introduces context manager pg.str_format and pg.repr_format.

These two (thread_local) context managers are used for controlling the behavior of __str__ and __repr__ for `pg.Formattable` objects.

PiperOrigin-RevId: 590276970

Author: daiyip

pyglove/core/object_utils/__init__.py

@@ -70,6 +70,7 @@
 """
 # pylint: enable=line-too-long
 # pylint: disable=g-bad-import-order
+# pylint: disable=g-importing-member
 
 # Common traits.
 from pyglove.core.object_utils.json_conversion import Nestable
@@ -116,6 +117,10 @@
 from pyglove.core.object_utils.formatting import BracketType
 from pyglove.core.object_utils.formatting import bracket_chars
 
+# Context managers for defining the default format for __str__ and __repr__.
+from pyglove.core.object_utils.common_traits import str_format
+from pyglove.core.object_utils.common_traits import repr_format
+
 # Handling code generation.
 from pyglove.core.object_utils.codegen import make_function
 
@@ -144,4 +149,5 @@
 from pyglove.core.object_utils.error_utils import catch_errors
 from pyglove.core.object_utils.error_utils import CatchErrorsContext
 
+# pylint: enable=g-importing-member
 # pylint: enable=g-bad-import-order

pyglove/core/object_utils/common_traits.py

@@ -19,7 +19,22 @@
 """
 
 import abc
-from typing import Any, Dict, Optional
+from typing import Any, ContextManager, Dict, Optional
+from pyglove.core.object_utils import thread_local
+
+
+_TLS_STR_FORMAT_KWARGS = '_str_format_kwargs'
+_TLS_REPR_FORMAT_KWARGS = '_repr_format_kwargs'
+
+
+def str_format(**kwargs) -> ContextManager[Dict[str, Any]]:
+  """Context manager for setting the default format kwargs for __str__."""
+  return thread_local.thread_local_arg_scope(_TLS_STR_FORMAT_KWARGS, **kwargs)
+
+
+def repr_format(**kwargs) -> ContextManager[Dict[str, Any]]:
+  """Context manager for setting the default format kwargs for __repr__."""
+  return thread_local.thread_local_arg_scope(_TLS_REPR_FORMAT_KWARGS, **kwargs)
 
 
 class Formattable(metaclass=abc.ABCMeta):
@@ -59,11 +74,15 @@ def format(self,
 
   def __str__(self) -> str:
     """Returns the full (maybe multi-line) representation of this object."""
-    return self.format(**self.__str_format_kwargs__)
+    kwargs = dict(self.__str_format_kwargs__)
+    kwargs.update(thread_local.thread_local_kwargs(_TLS_STR_FORMAT_KWARGS))
+    return self.format(**kwargs)
 
   def __repr__(self) -> str:
     """Returns a single-line representation of this object."""
-    return self.format(**self.__repr_format_kwargs__)
+    kwargs = dict(self.__repr_format_kwargs__)
+    kwargs.update(thread_local.thread_local_kwargs(_TLS_REPR_FORMAT_KWARGS))
+    return self.format(**kwargs)
 
 
 class MaybePartial(metaclass=abc.ABCMeta):
@@ -156,4 +175,3 @@ def ensure_explicit_method_override(
           f'{method} is a PyGlove managed method. If you do need to override '
           'it, please decorate the method with `@pg.explicit_method_override`.')
     raise TypeError(error_message)
-

pyglove/core/object_utils/common_traits_test.py

@@ -17,6 +17,36 @@
 from pyglove.core.object_utils import common_traits
 
 
+class Foo(common_traits.Formattable):
+
+  def format(self, compact: bool = False, verbose: bool = True, **kwargs):
+    return f'{self.__class__.__name__}(compact={compact}, verbose={verbose})'
+
+
+class FormattableTest(unittest.TestCase):
+
+  def test_formattable(self):
+    foo = Foo()
+    self.assertEqual(repr(foo), 'Foo(compact=True, verbose=True)')
+    self.assertEqual(str(foo), 'Foo(compact=False, verbose=True)')
+
+  def test_formattable_with_custom_format(self):
+    class Bar(Foo):
+      __str_format_kwargs__ = {'compact': False, 'verbose': False}
+      __repr_format_kwargs__ = {'compact': True, 'verbose': False}
+
+    bar = Bar()
+    self.assertEqual(repr(bar), 'Bar(compact=True, verbose=False)')
+    self.assertEqual(str(bar), 'Bar(compact=False, verbose=False)')
+
+  def test_formattable_with_context_managers(self):
+    foo = Foo()
+    with common_traits.str_format(verbose=False):
+      with common_traits.repr_format(compact=False):
+        self.assertEqual(repr(foo), 'Foo(compact=False, verbose=True)')
+        self.assertEqual(str(foo), 'Foo(compact=False, verbose=False)')
+
+
 class ExplicitlyOverrideTest(unittest.TestCase):
 
   def test_explicitly_override(self):

pyglove/core/object_utils/thread_local.py

@@ -15,11 +15,10 @@
 
 import contextlib
 import threading
-from typing import Any, Callable, Iterator
+from typing import Any, Callable, Dict, Iterator
 
-from pyglove.core.object_utils.missing import MISSING_VALUE
-
-_RAISE_IF_NOT_FOUND = (MISSING_VALUE,)
+_MISSING = KeyError()
+_RAISE_IF_NOT_FOUND = ValueError()
 _thread_local_state = threading.local()
 
 
@@ -40,6 +39,25 @@ def thread_local_value_scope(
       thread_local_del(key)
 
 
+@contextlib.contextmanager
+def thread_local_arg_scope(key: str, **kwargs) -> Iterator[Dict[str, Any]]:
+  """Context manager to update args associated with key."""
+  previous_kwargs = thread_local_peek(key, {})
+  current_kwargs = previous_kwargs.copy()
+  current_kwargs.update(kwargs)
+
+  try:
+    thread_local_push(key, current_kwargs)
+    yield current_kwargs
+  finally:
+    thread_local_pop(key)
+
+
+def thread_local_kwargs(key: str) -> Dict[str, Any]:
+  """Returns the args associated with key in current thread."""
+  return thread_local_peek(key, {})
+
+
 def thread_local_has(key: str) -> bool:
   """Deletes thread-local value by key."""
   return hasattr(_thread_local_state, key)
@@ -69,8 +87,8 @@ def thread_local_map(
     value_fn: Callable[[Any], Any],
     default_initial_value: Any = _RAISE_IF_NOT_FOUND) -> Any:
   """Map a thread-local value."""
-  value = thread_local_get(key, MISSING_VALUE)
-  if value == MISSING_VALUE:
+  value = thread_local_get(key, _MISSING)
+  if value is _MISSING:
     value = default_initial_value
     if value is _RAISE_IF_NOT_FOUND:
       raise ValueError(f'Key {key!r} does not exist in thread-local storage.')
@@ -112,10 +130,25 @@ def thread_local_push(key: str, value: Any) -> None:
   )
 
 
+def thread_local_peek(
+    key: str, default_value: Any = _RAISE_IF_NOT_FOUND
+) -> Any:
+  """Peaks a value at stack top."""
+  stack = thread_local_get(key, _MISSING)
+  if stack is _MISSING or not stack:
+    if default_value is _RAISE_IF_NOT_FOUND:
+      raise ValueError(
+          f'Stack associated with key {key!r} does not exist in thread-local '
+          'storage or is empty.'
+      )
+    return default_value
+  return stack[-1]
+
+
 def thread_local_pop(key: str, default_value: Any = _RAISE_IF_NOT_FOUND) -> Any:
   """Pops a value from a stack identified by key."""
-  stack = thread_local_get(key, MISSING_VALUE)
-  if stack == MISSING_VALUE:
+  stack = thread_local_get(key, _MISSING)
+  if stack is _MISSING:
     if default_value is _RAISE_IF_NOT_FOUND:
       raise ValueError(f'Key {key!r} does not exist in thread-local storage.')
     return default_value

pyglove/core/object_utils/thread_local_test.py

@@ -84,6 +84,33 @@ def _fn():
       return _fn
     self.assert_thread_func([thread_fun(i) for i in range(5)], 2)
 
+  def test_thread_local_arg_scope(self):
+    with thread_local.thread_local_arg_scope('arg_scope', x=1, y=2):
+      self.assertEqual(
+          thread_local.thread_local_kwargs('arg_scope'), dict(x=1, y=2)
+      )
+      with thread_local.thread_local_arg_scope('arg_scope', y=3, z=4):
+        self.assertEqual(
+            thread_local.thread_local_kwargs('arg_scope'), dict(x=1, y=3, z=4)
+        )
+      self.assertEqual(
+          thread_local.thread_local_kwargs('arg_scope'), dict(x=1, y=2)
+      )
+    self.assertEqual(thread_local.thread_local_kwargs('arg_scope'), dict())
+
+    # Test thread locality.
+    def thread_fun(i):
+      def _fn():
+        with thread_local.thread_local_arg_scope('arg_scope', x=i):
+          self.assertEqual(
+              thread_local.thread_local_kwargs('arg_scope'), dict(x=i)
+          )
+        self.assertEqual(thread_local.thread_local_kwargs('arg_scope'), dict())
+
+      return _fn
+
+    self.assert_thread_func([thread_fun(i) for i in range(5)], 2)
+
   def test_thread_local_increment_decrement(self):
     k = 'z'
     self.assertEqual(thread_local.thread_local_increment(k, 5), 6)
@@ -103,18 +130,29 @@ def _fn():
       return _fn
     self.assert_thread_func([thread_fun(i) for i in range(5)], 2)
 
-  def test_thread_local_push_pop(self):
+  def test_thread_local_push_peak_pop(self):
     k = 'p'
     self.assertFalse(thread_local.thread_local_has(k))
     thread_local.thread_local_push(k, 1)
     self.assertEqual(thread_local.thread_local_get(k), [1])
+    self.assertEqual(thread_local.thread_local_peek(k), 1)
     thread_local.thread_local_push(k, 2)
     self.assertEqual(thread_local.thread_local_get(k), [1, 2])
+    self.assertEqual(thread_local.thread_local_peek(k), 2)
     self.assertEqual(thread_local.thread_local_pop(k), 2)
     self.assertEqual(thread_local.thread_local_get(k), [1])
+    self.assertEqual(thread_local.thread_local_peek(k), 1)
     self.assertEqual(thread_local.thread_local_pop(k), 1)
+
+    with self.assertRaisesRegex(
+        ValueError, 'Stack associated with key .* does not exist'
+    ):
+      thread_local.thread_local_peek(k)
+    self.assertEqual(thread_local.thread_local_peek(k, -1), -1)
+
     with self.assertRaisesRegex(IndexError, 'pop from empty list'):
       thread_local.thread_local_pop(k)
+
     self.assertEqual(thread_local.thread_local_pop(k, -1), -1)
     with self.assertRaisesRegex(
         ValueError, 'Key .* does not exist in thread-local storage'):