pyglove 0.4.5.dev20240319__py3-none-any.whl → 0.4.5.dev202501140808__py3-none-any.whl

pyglove/core/symbolic/origin.py

@@ -16,11 +16,11 @@
 import traceback
 from typing import Any, Callable, List, Optional
 
-from pyglove.core import object_utils
+from pyglove.core import utils
 from pyglove.core.symbolic import flags
 
 
-class Origin(object_utils.Formattable):
+class Origin(utils.Formattable):
   """Class that represents the origin of a symbolic value.
 
   Origin is used for debugging the creation chain of a symbolic value, as
@@ -152,23 +152,26 @@ class Origin(object_utils.Formattable):
       compact: bool = False,
       verbose: bool = True,
       root_indent: int = 0,
-      *,
-      markdown: bool = False,
       **kwargs,
   ) -> str:
     """Formats this object."""
     if isinstance(self._source, (str, type(None))):
-      source_str = object_utils.quote_if_str(self._source)
+      source_str = self._source
     else:
-      source_info = object_utils.format(
-          self._source, compact, verbose, root_indent + 1, **kwargs)
-      source_str = f'{source_info} at 0x{id(self._source):8x}'
-    details = object_utils.kvlist_str([
-        ('tag', object_utils.quote_if_str(self._tag), None),
-        ('source', source_str, None),
-    ])
-    return object_utils.maybe_markdown_quote(
-        f'{self.__class__.__name__}({details})', markdown
+      source_info = utils.format(
+          self._source, compact, verbose, root_indent + 1, **kwargs
+      )
+      source_str = utils.RawText(f'{source_info} at 0x{id(self._source):8x}')
+
+    return utils.kvlist_str(
+        [
+            ('tag', self._tag, None),
+            ('source', source_str, None),
+        ],
+        label=self.__class__.__name__,
+        compact=compact,
+        verbose=verbose,
+        root_indent=root_indent,
     )
 
   def __eq__(self, other: Any) -> bool:

pyglove/core/symbolic/origin_test.py

@@ -92,11 +92,13 @@ class OriginTest(unittest.TestCase):
   def test_format(self):
     a = Dict(a=1)
     o = Origin(None, '__init__')
-    self.assertEqual(o.format(), 'Origin(tag=\'__init__\')')
+    self.assertEqual(o.format(compact=True), 'Origin(tag=\'__init__\')')
 
     o = Origin('/path/to/file', 'load')
     self.assertEqual(
-        o.format(), 'Origin(tag=\'load\', source=\'/path/to/file\')')
+        o.format(compact=False),
+        "Origin(\n  tag='load',\n  source='/path/to/file'\n)"
+    )
 
     o = Origin(a, 'builder')
     self.assertEqual(

pyglove/core/symbolic/pure_symbolic.py

@@ -14,8 +14,8 @@
 """Interfaces for pure symbolic objects."""
 
 from typing import Any, Callable, Optional, Tuple
-from pyglove.core import object_utils
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 
 
 class PureSymbolic(pg_typing.CustomTyping):
@@ -37,11 +37,12 @@ class PureSymbolic(pg_typing.CustomTyping):
 
   def custom_apply(
       self,
-      path: object_utils.KeyPath,
+      path: utils.KeyPath,
       value_spec: pg_typing.ValueSpec,
       allow_partial: bool,
       child_transform: Optional[
-          Callable[[object_utils.KeyPath, pg_typing.Field, Any], Any]] = None
+          Callable[[utils.KeyPath, pg_typing.Field, Any], Any]
+      ] = None,
   ) -> Tuple[bool, Any]:
     """Custom apply on a value based on its original value spec.
 

pyglove/core/symbolic/ref.py

@@ -13,15 +13,31 @@
 # limitations under the License.
 """Symbolic reference."""
 
-import numbers
-from typing import Any, Callable, Optional, Tuple
-from pyglove.core import object_utils
+import functools
+import typing
+from typing import Any, Callable, List, Optional, Tuple, Type
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.symbolic import base
-from pyglove.core.symbolic.object import Object
+from pyglove.core.symbolic import object as pg_object
+from pyglove.core.views.html import tree_view
 
 
-class Ref(Object, base.Inferential):
+class RefMeta(pg_object.ObjectMeta):
+  """Metaclass for Ref."""
+
+  def __getitem__(cls, type_arg: Type[Any]) -> Any:    # pylint: disable=no-self-argument
+    if typing.TYPE_CHECKING:
+      return type_arg
+    return pg_typing.Object(type_arg, transform=Ref)
+
+
+class Ref(
+    pg_object.Object,
+    base.Inferential,
+    tree_view.HtmlTreeView.Extension,
+    metaclass=RefMeta
+):
   """Symbolic reference.
 
   When adding a symbolic node to a symbolic tree, it undergoes a copy operation
@@ -80,11 +96,11 @@ class Ref(Object, base.Inferential):
 
   def __new__(cls, value: Any, **kwargs):
     del kwargs
-    if isinstance(value, (bool, numbers.Number, str)):
-      return value
-    return object.__new__(cls)
+    if isinstance(value, (base.Symbolic, list, dict)):
+      return object.__new__(cls)
+    return value
 
-  @object_utils.explicit_method_override
+  @utils.explicit_method_override
   def __init__(self, value: Any, **kwargs) -> None:
     super().__init__(**kwargs)
     if isinstance(value, Ref):
@@ -111,12 +127,13 @@ class Ref(Object, base.Inferential):
 
   def custom_apply(
       self,
-      path: object_utils.KeyPath,
+      path: utils.KeyPath,
       value_spec: pg_typing.ValueSpec,
       allow_partial: bool = False,
-      child_transform: Optional[Callable[
-          [object_utils.KeyPath, pg_typing.Field, Any], Any]] = None
-      ) -> Tuple[bool, Any]:
+      child_transform: Optional[
+          Callable[[utils.KeyPath, pg_typing.Field, Any], Any]
+      ] = None,
+  ) -> Tuple[bool, Any]:
     """Validate candidates during value_spec binding time."""
     del child_transform
     # Check if the field being assigned could accept the referenced value.
@@ -135,7 +152,9 @@ class Ref(Object, base.Inferential):
   def sym_eq(self, other: Any) -> bool:
     return isinstance(other, Ref) and self.value is other.value
 
-  def sym_jsonify(self, **kwargs: Any) -> Any:
+  def sym_jsonify(self, *, save_ref_value: bool = False, **kwargs: Any) -> Any:
+    if save_ref_value:
+      return base.to_json(self._value, save_ref_value=save_ref_value, **kwargs)
     raise TypeError(f'{self!r} cannot be serialized at the moment.')
 
   def __getstate__(self):
@@ -146,24 +165,68 @@ class Ref(Object, base.Inferential):
       compact: bool = False,
       verbose: bool = False,
       root_indent: int = 0,
-      *,
-      markdown: bool = False,
       **kwargs: Any,
   ) -> str:
-    value_str = object_utils.format(
+    value_str = utils.format(
         self._value,
-        compact=compact, verbose=verbose, root_indent=root_indent + 1)
+        compact=compact,
+        verbose=verbose,
+        root_indent=root_indent + 1,
+    )
     if compact:
-      s = f'{self.__class__.__name__}({value_str})'
+      return f'{self.__class__.__name__}({value_str})'
     else:
-      s = (
+      return (
           f'{self.__class__.__name__}(\n'
           + '  ' * (root_indent + 1)
           + f'value = {value_str}\n'
           + '  ' * root_indent
           + ')'
       )
-    return object_utils.maybe_markdown_quote(s, markdown)
+
+  def _html_tree_view_content(
+      self,
+      *,
+      view: tree_view.HtmlTreeView,
+      **kwargs: Any) -> tree_view.Html:
+    """Overrides `_html_content` to render the referenced value."""
+    return view.content(self._value, **kwargs)
+
+  def _html_tree_view_summary(
+      self,
+      *,
+      view: tree_view.HtmlTreeView,
+      title: Optional[str] = None,
+      **kwargs: Any) -> Optional[tree_view.Html]:
+    """Overrides `_html_content` to render the referenced value."""
+    return view.summary(
+        self,
+        title=title or f'{type(self._value).__name__}(...)',
+        **kwargs
+    )
+
+  @classmethod
+  @functools.cache
+  def _html_tree_view_config(cls) -> dict[str, Any]:
+    return tree_view.HtmlTreeView.get_kwargs(
+        super()._html_tree_view_config(),
+        dict(
+            css_classes=['ref'],
+        )
+    )
+
+  @classmethod
+  @functools.cache
+  def _html_tree_view_css_styles(cls) -> List[str]:
+    return super()._html_tree_view_css_styles() + [
+        """
+        /* Ref styles. */
+        .ref.summary-title::before {
+          content: 'ref: ';
+          color: #aaa;
+        }
+        """
+    ]
 
 
 def maybe_ref(value: Any) -> Optional[Ref]:
@@ -172,3 +235,27 @@ def maybe_ref(value: Any) -> Optional[Ref]:
     if value.sym_parent is None:
       return value
   return Ref(value)
+
+
+def deref(value: base.Symbolic, recursive: bool = False) -> Any:
+  """Dereferences a symbolic value that may contain pg.Ref.
+
+  Args:
+    value: The input symbolic value.
+    recursive: If True, dereference `pg.Ref` in the entire tree. Otherwise
+      Only dereference the root node.
+
+  Returns:
+    The dereferenced root, or dereferenced tree if recursive is True.
+  """
+  if isinstance(value, Ref):
+    value = value.value
+
+  if recursive:
+    def _deref(k, v, p):
+      del k, p
+      if isinstance(v, Ref):
+        return deref(v.value, recursive=True)
+      return v
+    return value.rebind(_deref, raise_on_no_change=False)
+  return value

pyglove/core/symbolic/ref_test.py

@@ -16,11 +16,13 @@
 import copy
 import inspect
 import pickle
+import typing
 from typing import Any
 import unittest
 
 from pyglove.core import typing as pg_typing
 from pyglove.core.symbolic import ref
+from pyglove.core.symbolic.base import contains
 from pyglove.core.symbolic.dict import Dict
 from pyglove.core.symbolic.object import Object
 
@@ -111,6 +113,11 @@ class RefTest(unittest.TestCase):
         TypeError, '.* cannot be serialized at the moment'):
       ref.Ref(A(1)).to_json()
 
+    self.assertEqual(
+        ref.Ref(A(1)).to_json(save_ref_value=True),
+        A(1).to_json()
+    )
+
   def test_pickle(self):
     with self.assertRaisesRegex(
         TypeError, '.* cannot be pickled at the moment'):
@@ -126,6 +133,92 @@ class RefTest(unittest.TestCase):
     self.assertIsInstance(x, ref.Ref)
     self.assertIs(x.value, a)
 
+  def test_deref(self):
+    class Foo(Object):
+      x: int
+      y: Any
+
+    # Deref on non-ref value.
+    a = Foo(1, 2)
+    self.assertIs(ref.deref(a), a)
+    self.assertEqual(ref.deref(a), Foo(1, 2))
+
+    # Deref top-level Ref.
+    a = ref.Ref(Foo(1, 2))
+    self.assertIs(ref.deref(a), a.value)
+
+    a = ref.Ref(Foo(1, ref.Ref(Foo(2, ref.Ref(Foo(3, 4))))))
+    self.assertTrue(contains(a, type=ref.Ref))
+    a_prime = ref.deref(a)
+    self.assertIs(a_prime, a.value)
+    self.assertIsInstance(a_prime.sym_getattr('y'), ref.Ref)
+    self.assertIsInstance(a_prime.y.sym_getattr('y'), ref.Ref)
+
+    # Deref entire tree.
+    a = ref.Ref(Foo(1, ref.Ref(Foo(2, ref.Ref(Foo(3, 4))))))
+    self.assertTrue(contains(a, type=ref.Ref))
+    a_prime = ref.deref(a, recursive=True)
+    self.assertIs(a_prime, a.value)
+    self.assertFalse(contains(a_prime, type=ref.Ref))
+
+  def test_to_html(self):
+
+    def assert_style(html, expected):
+      expected = inspect.cleandoc(expected).strip()
+      actual = html.style_section.strip()
+      if expected not in actual:
+        print(actual)
+      self.assertIn(expected, actual)
+
+    def assert_content(html, expected):
+      expected = inspect.cleandoc(expected).strip()
+      actual = html.content.strip()
+      if actual != expected:
+        print(actual)
+      self.assertEqual(actual.strip(), expected)
+
+    class Foo(Object):
+      x: Any
+
+    assert_style(
+        Foo(ref.Ref(Foo(1))).to_html(),
+        """
+        /* Ref styles. */
+        .ref.summary-title::before {
+          content: 'ref: ';
+          color: #aaa;
+        }
+        """
+    )
+    assert_content(
+        Foo(ref.Ref(Foo(1))).to_html(
+            extra_flags=dict(
+                use_inferred=False,
+            ),
+            enable_summary_tooltip=False,
+            enable_key_tooltip=False,
+        ),
+        """
+        <details open class="pyglove foo"><summary><div class="summary-title">Foo(...)</div></summary><div class="complex-value foo"><details class="pyglove ref"><summary><div class="summary-name ref">x</div><div class="summary-title ref">Foo(...)</div></summary><div class="complex-value foo"><details open class="pyglove int"><summary><div class="summary-name">x</div><div class="summary-title">int</div></summary><span class="simple-value int">1</span></details></div></details></div></details>
+        """
+    )
+
+  def test_annotation(self):
+    typing.TYPE_CHECKING = True
+    assert ref.Ref[ValueError] is ValueError
+    typing.TYPE_CHECKING = False
+
+    class Bar(Object):
+      x: int
+
+    class Foo(Object):
+      y: ref.Ref[Bar]
+
+    f = Foo(Bar(1))
+    self.assertIsInstance(f.sym_getattr('y'), ref.Ref)
+    self.assertEqual(f.y.sym_path, '')
+    self.assertIsInstance(ref.Ref[Bar](1), ref.Ref)
+
 
 if __name__ == '__main__':
   unittest.main()

pyglove/core/symbolic/symbolize_test.py

@@ -73,8 +73,16 @@ class SymbolizeFunctionsTest(unittest.TestCase):
     self.assertEqual(
         f.__signature__.args,
         [
-            (pg_typing.Argument('x', pg_typing.Int())),
-            (pg_typing.Argument('y', pg_typing.Str())),
+            pg_typing.Argument(
+                'x',
+                pg_typing.Argument.Kind.POSITIONAL_OR_KEYWORD,
+                pg_typing.Int()
+            ),
+            pg_typing.Argument(
+                'y',
+                pg_typing.Argument.Kind.POSITIONAL_OR_KEYWORD,
+                pg_typing.Str()
+            ),
         ],
     )
     self.assertEqual(f.__signature__.return_value, pg_typing.Int())

pyglove/core/tuning/local_backend.py

@@ -21,8 +21,8 @@ from typing import Any, Callable, Dict, List, Optional, Sequence
 
 from pyglove.core import geno
 from pyglove.core import logging
-from pyglove.core import object_utils
 from pyglove.core import symbolic
+from pyglove.core import utils
 from pyglove.core.tuning import backend
 from pyglove.core.tuning.early_stopping import EarlyStoppingPolicy
 from pyglove.core.tuning.protocols import Feedback
@@ -278,7 +278,7 @@ class _InMemoryResult(Result):
           ('step', self._best_trial.final_measurement.step),
           ('dna', self._best_trial.dna.format(compact=True))
       ])
-    return object_utils.format(json_repr, compact, False, root_indent, **kwargs)
+    return utils.format(json_repr, compact, False, root_indent, **kwargs)
 
 
 @backend.add_backend('in-memory')

pyglove/core/tuning/protocols.py

@@ -22,9 +22,9 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple, Type, Union
 
 from pyglove.core import geno
 from pyglove.core import logging
-from pyglove.core import object_utils
 from pyglove.core import symbolic
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 
 
 class _DataEntity(symbolic.Object):
@@ -116,7 +116,7 @@ class Trial(_DataEntity):
     return tuple(metric_values) if len(metric_values) > 1 else metric_values[0]
 
 
-class Result(object_utils.Formattable):
+class Result(utils.Formattable):
   """Interface for tuning result."""
 
   @property
@@ -416,7 +416,7 @@ class Feedback(metaclass=abc.ABCMeta):
       error_stack = traceback.format_exc()
       logging.warning('Skipping trial on unhandled exception: %s', error_stack)
       self.skip(error_stack)
-    return object_utils.catch_errors(exceptions, skip_on_exception)
+    return utils.catch_errors(exceptions, skip_on_exception)
 
   @contextlib.contextmanager
   def ignore_race_condition(self):

pyglove/core/tuning/sample_test.py

@@ -173,15 +173,15 @@ class SamplingTest(unittest.TestCase):
     _, f = next(sample)
     with self.assertRaisesRegex(
         ValueError, 'bcd'):
-      with f.skip_on_exceptions(((Exception, '.*a'),)):
+      with f.skip_on_exceptions(((Exception, 'abc'),)):
         # should skip.
         raise ValueError('bcd')
     self.assertEqual(algo.num_proposals, 7)
     self.assertEqual(algo.num_feedbacks, 1)
 
     _, f = next(sample)
-    with f.skip_on_exceptions(((ValueError, '.*a'),
-                               (ValueError, '.*b'),
+    with f.skip_on_exceptions(((ValueError, 'abc'),
+                               (ValueError, 'cd'),
                                KeyError)):
       # should skip.
       raise ValueError('bcd')

pyglove/core/typing/__init__.py

@@ -303,15 +303,14 @@ from pyglove.core.typing.typed_missing import MissingValue   # Typed.
 from pyglove.core.typing.class_schema import KeySpec
 from pyglove.core.typing.class_schema import ValueSpec
 from pyglove.core.typing.class_schema import Field
+from pyglove.core.typing.class_schema import FieldKeyDef
+from pyglove.core.typing.class_schema import FieldValueDef
+from pyglove.core.typing.class_schema import FieldDef
 from pyglove.core.typing.class_schema import Schema
 from pyglove.core.typing.class_schema import create_field
 from pyglove.core.typing.class_schema import create_schema
 from pyglove.core.typing.class_schema import ForwardRef
 
-# Class schema helpers.
-from pyglove.core.typing.class_schema_utils import get_arg_fields
-from pyglove.core.typing.class_schema_utils import ensure_value_spec
-
 # Concrete key specifications.
 from pyglove.core.typing.key_specs import ConstStrKey
 from pyglove.core.typing.key_specs import NonConstKey
@@ -337,6 +336,8 @@ from pyglove.core.typing.value_specs import Type
 from pyglove.core.typing.value_specs import Union
 from pyglove.core.typing.value_specs import Any
 
+from pyglove.core.typing.value_specs import ensure_value_spec
+
 # Generic type aliases.
 from pyglove.core.typing.value_specs import GenericMeta
 from pyglove.core.typing.value_specs import Generic
@@ -361,6 +362,7 @@ from pyglove.core.typing.type_conversion import get_json_value_converter
 # Inspect helpers.
 from pyglove.core.typing.inspect import is_subclass
 from pyglove.core.typing.inspect import is_instance
+from pyglove.core.typing.inspect import get_outer_class
 from pyglove.core.typing.inspect import get_type
 from pyglove.core.typing.inspect import get_type_args
 from pyglove.core.typing.inspect import is_generic
@@ -377,9 +379,16 @@ from pyglove.core.typing.custom_typing import CustomTyping
 from pyglove.core.typing.callable_signature import Argument
 from pyglove.core.typing.callable_signature import CallableType
 from pyglove.core.typing.callable_signature import Signature
-from pyglove.core.typing.callable_signature import get_signature
+from pyglove.core.typing.callable_signature import signature
+from pyglove.core.typing.callable_signature import schema
+
+# For backward compatibility.
+get_signature = signature
 
 # Callable extensions.
+from pyglove.core.typing.callable_ext import PresetArgValue
+from pyglove.core.typing.callable_ext import enable_preset_args
+from pyglove.core.typing.callable_ext import preset_args
 from pyglove.core.typing.callable_ext import CallableWithOptionalKeywordArgs
 
 # PyType support.