pyglove 0.4.5.dev20240318__py3-none-any.whl → 0.4.5.dev202501132210__py3-none-any.whl

pyglove/core/utils/__init__.py

@@ -0,0 +1,164 @@
+# Copyright 2022 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.
+# pylint: disable=line-too-long
+"""Utility library that provides common traits for objects in Python.
+
+Overview
+--------
+
+``pg.utils`` sits at the bottom of all PyGlove modules and empowers other
+modules with the following features:
+
+  +---------------------+--------------------------------------------+
+  | Functionality       | API                                        |
+  +=====================+============================================+
+  | Formatting          | :class:`pg.Formattable`,                   |
+  |                     |                                            |
+  |                     | :func:`pg.format`,                         |
+  |                     |                                            |
+  |                     | :func:`pg.print`,                          |
+  |                     |                                            |
+  |                     | :func:`pg.utils.kvlist_str`,        |
+  |                     |                                            |
+  |                     | :func:`pg.utils.quote_if_str`,      |
+  |                     |                                            |
+  |                     | :func:`pg.utils.message_on_path`    |
+  +---------------------+--------------------------------------------+
+  | Serialization       | :class:`pg.JSONConvertible`,               |
+  |                     |                                            |
+  |                     | :func:`pg.registered_types`,               |
+  |                     |                                            |
+  |                     | :func:`pg.utils.to_json`,           |
+  |                     |                                            |
+  |                     | :func:`pg.utils.from_json`,         |
+  +---------------------+--------------------------------------------+
+  | Partial construction| :class:`pg.MaybePartial`,                  |
+  |                     |                                            |
+  |                     | :const:`pg.MISSING_VALUE`.                 |
+  +---------------------+--------------------------------------------+
+  | Hierarchical key    | :class:`pg.KeyPath`                        |
+  | representation      |                                            |
+  +---------------------+--------------------------------------------+
+  | Hierarchical object | :func:`pg.utils.traverse`           |
+  | traversal           |                                            |
+  +---------------------+--------------------------------------------+
+  | Hierarchical object | :func:`pg.utils.transform`,         |
+  | transformation      |                                            |
+  |                     | :func:`pg.utils.merge`,             |
+  |                     |                                            |
+  |                     | :func:`pg.utils.canonicalize`,      |
+  |                     |                                            |
+  |                     | :func:`pg.utils.flatten`            |
+  +---------------------+--------------------------------------------+
+  | Docstr handling     | :class:`pg.docstr`,                        |
+  +---------------------+--------------------------------------------+
+  | Error handling     | :class:`pg.catch_errors`,                   |
+  +---------------------+--------------------------------------------+
+"""
+# pylint: enable=line-too-long
+# pylint: disable=g-bad-import-order
+# pylint: disable=g-importing-member
+
+# Handling JSON conversion.
+from pyglove.core.utils.json_conversion import Nestable
+from pyglove.core.utils.json_conversion import JSONValueType
+
+from pyglove.core.utils.json_conversion import JSONConvertible
+from pyglove.core.utils.json_conversion import from_json
+from pyglove.core.utils.json_conversion import to_json
+from pyglove.core.utils.json_conversion import registered_types
+
+# Handling formatting.
+from pyglove.core.utils.formatting import Formattable
+from pyglove.core.utils.formatting import format  # pylint: disable=redefined-builtin
+from pyglove.core.utils.formatting import printv as print  # pylint: disable=redefined-builtin
+from pyglove.core.utils.formatting import kvlist_str
+from pyglove.core.utils.formatting import quote_if_str
+from pyglove.core.utils.formatting import maybe_markdown_quote
+from pyglove.core.utils.formatting import comma_delimited_str
+from pyglove.core.utils.formatting import camel_to_snake
+from pyglove.core.utils.formatting import auto_plural
+from pyglove.core.utils.formatting import BracketType
+from pyglove.core.utils.formatting import bracket_chars
+from pyglove.core.utils.formatting import RawText
+
+# Context managers for defining the default format for __str__ and __repr__.
+from pyglove.core.utils.formatting import str_format
+from pyglove.core.utils.formatting import repr_format
+
+# Value location.
+from pyglove.core.utils.value_location import KeyPath
+from pyglove.core.utils.value_location import KeyPathSet
+from pyglove.core.utils.value_location import StrKey
+from pyglove.core.utils.value_location import message_on_path
+
+# Value markers.
+from pyglove.core.utils.missing import MissingValue
+from pyglove.core.utils.missing import MISSING_VALUE
+
+# Handling hierarchical.
+from pyglove.core.utils.hierarchical import traverse
+from pyglove.core.utils.hierarchical import transform
+from pyglove.core.utils.hierarchical import flatten
+from pyglove.core.utils.hierarchical import canonicalize
+from pyglove.core.utils.hierarchical import merge
+from pyglove.core.utils.hierarchical import merge_tree
+from pyglove.core.utils.hierarchical import is_partial
+from pyglove.core.utils.hierarchical import try_listify_dict_with_int_keys
+
+# Common traits.
+from pyglove.core.utils.common_traits import MaybePartial
+from pyglove.core.utils.common_traits import Functor
+
+from pyglove.core.utils.common_traits import explicit_method_override
+from pyglove.core.utils.common_traits import ensure_explicit_method_override
+
+# Handling thread local values.
+from pyglove.core.utils.thread_local import thread_local_value_scope
+from pyglove.core.utils.thread_local import thread_local_has
+from pyglove.core.utils.thread_local import thread_local_set
+from pyglove.core.utils.thread_local import thread_local_get
+from pyglove.core.utils.thread_local import thread_local_del
+from pyglove.core.utils.thread_local import thread_local_increment
+from pyglove.core.utils.thread_local import thread_local_decrement
+from pyglove.core.utils.thread_local import thread_local_push
+from pyglove.core.utils.thread_local import thread_local_pop
+from pyglove.core.utils.thread_local import thread_local_peek
+
+# Handling docstrings.
+from pyglove.core.utils.docstr_utils import DocStr
+from pyglove.core.utils.docstr_utils import DocStrStyle
+from pyglove.core.utils.docstr_utils import DocStrEntry
+from pyglove.core.utils.docstr_utils import DocStrExample
+from pyglove.core.utils.docstr_utils import DocStrArgument
+from pyglove.core.utils.docstr_utils import DocStrReturns
+from pyglove.core.utils.docstr_utils import DocStrRaises
+from pyglove.core.utils.docstr_utils import docstr
+
+# Handling exceptions.
+from pyglove.core.utils.error_utils import catch_errors
+from pyglove.core.utils.error_utils import CatchErrorsContext
+from pyglove.core.utils.error_utils import ErrorInfo
+
+# Timing.
+from pyglove.core.utils.timing import timeit
+from pyglove.core.utils.timing import TimeIt
+
+# Text color.
+from pyglove.core.utils.text_color import colored
+from pyglove.core.utils.text_color import colored_block
+from pyglove.core.utils.text_color import decolor
+
+# pylint: enable=g-importing-member
+# pylint: enable=g-bad-import-order

pyglove/core/{object_utils → utils}/common_traits.py

@@ -14,75 +14,11 @@
 """Common traits for Python objects.
 
 This file defines interfaces for describing the common traits of a Python
-object, for example, partiality (MaybePartial), formatting (Formattable),
-functor (Functor).
+object, for example, partiality (MaybePartial), functor (Functor).
 """
 
 import abc
-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):
-  """Interface for classes whose instances can be pretty-formatted.
-
-  This interface overrides the default ``__repr__`` and ``__str__`` method, thus
-  all ``Formattable`` objects can be printed nicely.
-
-  All symbolic types implement this interface.
-  """
-
-  # Additional format keyword arguments for `__str__`.
-  __str_format_kwargs__ = dict(compact=False, verbose=True)
-
-  # Additional format keyword arguments for `__repr__`.
-  __repr_format_kwargs__ = dict(compact=True)
-
-  @abc.abstractmethod
-  def format(self,
-             compact: bool = False,
-             verbose: bool = True,
-             root_indent: int = 0,
-             **kwargs) -> str:
-    """Formats this object into a string representation.
-
-    Args:
-      compact: If True, this object will be formatted into a single line.
-      verbose: If True, this object will be formatted with verbosity.
-        Subclasses should define `verbosity` on their own.
-      root_indent: The start indent level for this object if the output is a
-        multi-line string.
-      **kwargs: Subclass specific keyword arguments.
-
-    Returns:
-      A string of formatted object.
-    """
-
-  def __str__(self) -> str:
-    """Returns the full (maybe multi-line) representation of this object."""
-    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."""
-    kwargs = dict(self.__repr_format_kwargs__)
-    kwargs.update(thread_local.thread_local_kwargs(_TLS_REPR_FORMAT_KWARGS))
-    return self.format(**kwargs)
+from typing import Any, Dict, Optional, Union
 
 
 class MaybePartial(metaclass=abc.ABCMeta):
@@ -111,7 +47,7 @@ class MaybePartial(metaclass=abc.ABCMeta):
     return len(self.missing_values()) > 0  # pylint: disable=g-explicit-length-test
 
   @abc.abstractmethod
-  def missing_values(self, flatten: bool = True) -> Dict[str, Any]:  # pylint: disable=redefined-outer-name
+  def missing_values(self, flatten: bool = True) -> Dict[Union[str, int], Any]:  # pylint: disable=redefined-outer-name
     """Returns missing values from this object.
 
     Args:

pyglove/core/utils/common_traits_test.py

@@ -0,0 +1,36 @@
+# Copyright 2022 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 unittest
+from pyglove.core.utils import common_traits
+
+
+class ExplicitlyOverrideTest(unittest.TestCase):
+
+  def test_explicitly_override(self):
+    class A:
+
+      @common_traits.explicit_method_override
+      def __init__(self, x, y):
+        pass
+
+      def bar(self):
+        pass
+
+    common_traits.ensure_explicit_method_override(A.__init__)
+    with self.assertRaisesRegex(TypeError, '.* is a PyGlove managed method'):
+      common_traits.ensure_explicit_method_override(A.bar)
+
+
+if __name__ == '__main__':
+  unittest.main()

pyglove/core/{object_utils → utils}/docstr_utils.py

@@ -15,6 +15,7 @@
 
 import dataclasses
 import enum
+import inspect
 from typing import Any, Dict, List, Optional
 
 import docstring_parser
@@ -74,6 +75,28 @@ class DocStr:
   raises: List[DocStrRaises]
   blank_after_short_description: bool = True
 
+  @property
+  def description(self) -> Optional[str]:
+    """Returns short_description + long_description."""
+    if self.short_description is None and self.long_description is None:
+      return None
+    description = self.short_description or ''
+    if self.blank_after_short_description:
+      description += '\n'
+    if self.long_description:
+      description += '\n' + self.long_description
+    return description.rstrip('\n')
+
+  def parameter(self, param: inspect.Parameter) -> Optional[DocStrArgument]:
+    """Returns doc str for an inspected parameter."""
+    if param.kind == inspect.Parameter.VAR_POSITIONAL:
+      name = f'*{param.name}'
+    elif param.kind == inspect.Parameter.VAR_KEYWORD:
+      name = f'**{param.name}'
+    else:
+      name = param.name
+    return self.args.get(name)
+
   @classmethod
   def parse(cls, text: str, style: Optional[DocStrStyle] = None) -> 'DocStr':
     """Parses a docstring."""

pyglove/core/{object_utils → utils}/docstr_utils_test.py

@@ -11,10 +11,9 @@
 # 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.object_utils.docstr_utils."""
-
+import inspect
 import unittest
-from pyglove.core.object_utils import docstr_utils
+from pyglove.core.utils import docstr_utils
 
 
 class DocStrTest(unittest.TestCase):
@@ -57,7 +56,16 @@ class DocStrTest(unittest.TestCase):
       del args, kwargs
       return x + y
 
-    self.assertEqual(docstr_utils.docstr(my_sum), docstr_utils.DocStr(
+    docstr = docstr_utils.docstr(my_sum)
+    self.assertEqual(
+        docstr.description,
+        inspect.cleandoc('''
+        Returns the sum of two integers.
+
+        This function will return the sum of two integers.
+        ''')
+    )
+    self.assertEqual(docstr, docstr_utils.DocStr(
         style=docstr_utils.DocStrStyle.GOOGLE,
         short_description='Returns the sum of two integers.',
         long_description='This function will return the sum of two integers.',
@@ -93,6 +101,11 @@ class DocStrTest(unittest.TestCase):
             )
         ]
     ))
+    sig = inspect.signature(my_sum)
+    self.assertIsNotNone(docstr.parameter(sig.parameters['x']))
+    self.assertIsNotNone(docstr.parameter(sig.parameters['y']))
+    self.assertIsNotNone(docstr.parameter(sig.parameters['args']))
+    self.assertIsNotNone(docstr.parameter(sig.parameters['kwargs']))
 
     class Foo:
       pass
@@ -100,6 +113,25 @@ class DocStrTest(unittest.TestCase):
     self.assertIsNone(docstr_utils.docstr(Foo))
     self.assertIsNone(docstr_utils.docstr(None))
 
+    class Bar:
+      """bar."""
+
+    self.assertEqual(docstr_utils.docstr(Bar).description, 'bar.')
+
+    # pylint: disable=g-classes-have-attributes
+    # pylint: disable=g-short-docstring-punctuation
+    # pylint: disable=g-space-before-docstring-summary
+    # pylint: disable=g-no-space-after-docstring-summary
+    class Baz:
+      """
+      Args:
+        x: int
+      """
+    # pylint: enable=g-space-before-docstring-summary
+    # pylint: enable=g-short-docstring-punctuation
+    # pylint: enable=g-classes-have-attributes
+    self.assertIsNone(docstr_utils.docstr(Baz).description)
+
 
 if __name__ == '__main__':
   unittest.main()

pyglove/core/{object_utils → utils}/error_utils.py

@@ -17,28 +17,92 @@ import contextlib
 import dataclasses
 import inspect
 import re
-from typing import Callable, Dict, List, Optional, Sequence, Tuple, Type, Union
+import sys
+import traceback
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type, Union
+
+from pyglove.core.utils import formatting
+from pyglove.core.utils import json_conversion
+
+
+@dataclasses.dataclass(frozen=True)
+class ErrorInfo(json_conversion.JSONConvertible, formatting.Formattable):
+  """Serializable error information.
+
+  Attributes:
+    tag: A path of the error types in the exception chain. For example,
+      `ValueError.ZeroDivisionError` means the error is a `ZeroDivisionError`
+      raised at the first place and then reraised as a `ValueError`.
+    description: The description of the error.
+    stacktrace: The stacktrace of the error.
+  """
+
+  tag: str
+  description: str
+  stacktrace: str
+
+  @classmethod
+  def _compute_tag(cls, error: BaseException):
+    error_types = []
+    while error is not None:
+      error_types.append(error.__class__.__name__)
+      error = getattr(error, 'cause', error.__cause__)
+    return '.'.join(error_types)
+
+  @classmethod
+  def from_exception(cls, error: BaseException) -> 'ErrorInfo':
+    """Creates an error info from an exception."""
+    return cls(
+        tag=cls._compute_tag(error),
+        description=str(error),
+        stacktrace=''.join(
+            traceback.format_exception(*sys.exc_info())
+        )
+    )
+
+  def to_json(self, **kwargs) -> Dict[str, Any]:
+    return self.to_json_dict(
+        fields=dict(
+            tag=(self.tag, None),
+            description=(self.description, None),
+            stacktrace=(self.stacktrace, None),
+        ),
+        exclude_default=True,
+        **kwargs,
+    )
+
+  def format(self, *args, **kwargs) -> str:
+    return formatting.kvlist_str(
+        [
+            ('tag', self.tag, None),
+            ('description', self.description, None),
+            ('stacktrace', self.stacktrace, None),
+        ],
+        *args,
+        label=self.__class__.__name__,
+        **kwargs,
+    )
 
 
 @dataclasses.dataclass()
 class CatchErrorsContext:
   """Context for pg.catch_errors."""
-  error: Optional[Exception] = None
+  error: Optional[BaseException] = None
 
 
 @contextlib.contextmanager
 def catch_errors(
     errors: Union[
-        Union[Type[Exception], Tuple[Exception, str]],
-        Sequence[Union[Type[Exception], Tuple[Exception, str]]],
+        Union[Type[BaseException], Tuple[Type[BaseException], str]],
+        Sequence[Union[Type[BaseException], Tuple[Type[BaseException], str]]],
     ],
-    error_handler: Optional[Callable[[Exception], None]] = None
+    error_handler: Optional[Callable[[BaseException], None]] = None
 ):
   """Context manager for catching user-specified exceptions.
 
   Examples::
 
-    with pg.object_utils.catch_errors(
+    with pg.utils.catch_errors(
         [
             RuntimeErrror,
             (ValueError, 'Input is wrong.')
@@ -71,7 +135,7 @@ def catch_errors(
   ):
     errors = [errors]
 
-  error_mapping: Dict[Type[Exception], List[str]] = {}
+  error_mapping: Dict[Type[BaseException], List[str]] = {}
   for error_type in errors:
     regex = None
     if isinstance(error_type, tuple):
@@ -82,7 +146,9 @@ def catch_errors(
             f'to match. Encountered: {error_type!r}.'
         )
       error_type, regex = error_type
-    if not (inspect.isclass(error_type) and issubclass(error_type, Exception)):
+    if not (
+        inspect.isclass(error_type) and issubclass(error_type, BaseException)
+    ):
       raise TypeError(f'Exception contains non-except types: {error_type!r}.')
     if error_type not in error_mapping:
       error_mapping[error_type] = []
@@ -93,7 +159,7 @@ def catch_errors(
   try:
     yield context
   except tuple(error_mapping.keys()) as e:
-    error_message = str(e)
+    error_message = e.__class__.__name__ + ': ' + str(e)
     found_match = False
     for error_type, error_regexes in error_mapping.items():
       if isinstance(e, error_type):
@@ -101,6 +167,9 @@ def catch_errors(
           found_match = True
         else:
           for regex in error_regexes:
+            assert regex is not None
+            if not regex.startswith(('^', '.*')):
+              regex = '.*' + regex
             if re.match(regex, error_message):
               found_match = True
               break

pyglove/core/{object_utils → utils}/error_utils_test.py

@@ -11,10 +11,64 @@
 # 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.object_utils.error_utils."""
-
+import inspect
 import unittest
-from pyglove.core.object_utils import error_utils
+from pyglove.core.utils import error_utils
+
+
+class ErrorInfoTest(unittest.TestCase):
+  """Tests for ErrorInfo."""
+
+  def test_from_exception(self):
+
+    def foo():
+      return 1 / 0
+
+    def bar():
+      try:
+        foo()
+      except ZeroDivisionError as e:
+        raise ValueError('Bad call to `foo`') from e
+
+    error_info = None
+    try:
+      bar()
+    except ValueError as e:
+      error_info = error_utils.ErrorInfo.from_exception(e)
+    self.assertIsNotNone(error_info)
+    self.assertEqual(error_info.tag, 'ValueError.ZeroDivisionError')
+    self.assertEqual(error_info.description, 'Bad call to `foo`')
+    self.assertIn('Traceback (most recent call last)', error_info.stacktrace)
+
+  def test_to_json(self):
+    error_info = error_utils.ErrorInfo(
+        tag='ValueError.ZeroDivisionError',
+        description='Bad call to `foo`',
+        stacktrace='Traceback (most recent call last)',
+    )
+    json_dict = error_info.to_json()
+    error_info2 = error_utils.ErrorInfo.from_json(json_dict)
+    self.assertIsNot(error_info2, error_info)
+    self.assertEqual(error_info2, error_info)
+
+  def test_format(self):
+    error_info = error_utils.ErrorInfo(
+        tag='ValueError.ZeroDivisionError',
+        description='Bad call to `foo`',
+        stacktrace='Traceback (most recent call last)',
+    )
+    self.assertEqual(
+        error_info.format(compact=False),
+        inspect.cleandoc(
+            """
+            ErrorInfo(
+              tag='ValueError.ZeroDivisionError',
+              description='Bad call to `foo`',
+              stacktrace='Traceback (most recent call last)'
+            )
+            """
+        )
+    )
 
 
 class CatchErrorsTest(unittest.TestCase):
@@ -36,11 +90,13 @@ class CatchErrorsTest(unittest.TestCase):
 
     self.assert_caught_error(foo, ValueError)
     self.assert_caught_error(foo, (ValueError,))
+    self.assert_caught_error(foo, (Exception, 'ValueError'))
     self.assert_caught_error(foo, (KeyError, ValueError))
-    self.assert_caught_error(foo, (ValueError, '.* an error'))
-    self.assert_caught_error(foo, (KeyError, (ValueError, '.* an error'),))
+    self.assert_caught_error(foo, (ValueError, 'an error'))
+    self.assert_caught_error(foo, (KeyError, (ValueError, 'an error'),))
 
     self.assert_propagate_error(foo, KeyError)
+    self.assert_propagate_error(foo, (ValueError, '^an error'))
     self.assert_propagate_error(foo, (ValueError, 'something else'))
 
   def test_catch_errors_with_error_handler(self):