pyglove 0.4.5.dev202410290809__py3-none-any.whl → 0.4.5.dev202410301815__py3-none-any.whl

pyglove/core/__init__.py

@@ -163,6 +163,7 @@ to_json = symbolic.to_json
 to_json_str = symbolic.to_json_str
 save = symbolic.save
 load = symbolic.load
+open_jsonl = symbolic.open_jsonl
 get_load_handler = symbolic.get_load_handler
 set_load_handler = symbolic.set_load_handler
 get_save_handler = symbolic.get_save_handler

pyglove/core/io/__init__.py

@@ -16,5 +16,6 @@
 # pylint: disable=g-bad-import-order
 
 from pyglove.core.io.file_system import *
+from pyglove.core.io.sequence import *
 
 # pylint: enable=g-bad-import-order

pyglove/core/io/file_system.py

@@ -51,6 +51,10 @@ class File(metaclass=abc.ABCMeta):
   def tell(self) -> int:
     """Returns the current position of the file."""
 
+  @abc.abstractmethod
+  def flush(self) -> None:
+    """Flushes the written content to the storage."""
+
   @abc.abstractmethod
   def close(self) -> None:
     """Closes the file."""
@@ -148,6 +152,9 @@ class StdFile(File):
   def tell(self) -> int:
     return self._file_object.tell()
 
+  def flush(self) -> None:
+    self._file_object.flush()
+
   def close(self) -> None:
     self._file_object.close()
 
@@ -220,6 +227,9 @@ class MemoryFile(File):
   def tell(self) -> int:
     return self._buffer.tell()
 
+  def flush(self) -> None:
+    pass
+
   def close(self) -> None:
     self.seek(0)
 

pyglove/core/io/file_system_test.py

@@ -133,6 +133,7 @@ class MemoryFileSystemTest(unittest.TestCase):
     self.assertFalse(fs.exists(file1))
     with fs.open(file1, 'w') as f:
       f.write('hello')
+      f.flush()
     self.assertTrue(fs.exists(file1))
     self.assertFalse(fs.isdir(file1))
 
@@ -224,6 +225,7 @@ class FileIoApiTest(unittest.TestCase):
     with file_system.open(file1, 'w') as f:
       self.assertIsInstance(f, file_system.MemoryFile)
       f.write('foo')
+      f.flush()
     self.assertTrue(file_system.path_exists(file1))
     self.assertEqual(file_system.readfile(file1), 'foo')
     file_system.writefile(file1, 'bar')

pyglove/core/io/sequence.py

@@ -0,0 +1,301 @@
+# Copyright 2024 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.
+"""Pluggable record IO."""
+
+import abc
+import collections
+import os
+from typing import Any, Callable, Iterator, Optional, Union
+
+from pyglove.core.io import file_system
+
+
+class Sequence(metaclass=abc.ABCMeta):
+  """Interface for a sequence of records."""
+
+  def __init__(
+      self,
+      serializer: Optional[Callable[[Any], Union[bytes, str]]] = None,
+      deserializer: Optional[Callable[[Union[bytes, str]], Any]] = None
+  ):
+    self._serializer = serializer
+    self._deserializer = deserializer
+
+  def add(self, record: Any) -> None:
+    """Adds a record to the reader."""
+    if self._serializer:
+      record = self._serializer(record)
+    if not isinstance(record, (str, bytes)):
+      raise ValueError(
+          f'Cannot write record with type {type(record)}. '
+          'Did you forget to pass a serializer?'
+      )
+    self._add(record)
+
+  @abc.abstractmethod
+  def _add(self, record: Union[bytes, str]) -> None:
+    """Adds a raw record to the reader."""
+
+  @abc.abstractmethod
+  def __len__(self):
+    """Gets the number of records in the reader."""
+
+  def __iter__(self) -> Iterator[Any]:
+    """Iterates over the records in the reader."""
+    for record in self._iter():
+      if self._deserializer:
+        yield self._deserializer(record)
+      else:
+        yield record
+
+  @abc.abstractmethod
+  def _iter(self) -> Iterator[Union[bytes, str]]:
+    """Iterates over the raw records in the reader."""
+
+  @abc.abstractmethod
+  def close(self) -> None:
+    """Closes the reader."""
+
+  @abc.abstractmethod
+  def flush(self) -> None:
+    """Flushes the read records to the storage."""
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    del exc_type, exc_value, traceback
+    self.close()
+
+
+class SequenceIO(metaclass=abc.ABCMeta):
+  """Interface for a record IO system."""
+
+  @abc.abstractmethod
+  def open(
+      self,
+      path: Union[str, os.PathLike[str]],
+      mode: str,
+      *,
+      serializer: Optional[Callable[[Any], Union[bytes, str]]],
+      deserializer: Optional[Callable[[Union[bytes, str]], Any]],
+      **kwargs
+  ) -> Sequence:
+    """Opens a sequence for reading or writing."""
+
+
+class _SequenceIORegistry(object):
+  """Registry for record IO systems."""
+
+  def __init__(self):
+    self._registry = {}
+
+  def add(self, extension: str, sequence_io: SequenceIO) -> None:
+    """Adds a record IO system with a prefix."""
+    self._registry[extension] = sequence_io
+
+  def get(self, path: Union[str, os.PathLike[str]]) -> SequenceIO:
+    """Gets the record IO system for a path."""
+    path = _resolve_path(path)
+    parts = path.split('.')
+    if parts:
+      extension = parts[-1].lower()
+      if '@' in extension:
+        extension = extension.split('@')[0]
+      if extension in self._registry:
+        return self._registry[extension]
+    return LineSequenceIO()
+
+
+_registry = _SequenceIORegistry()
+
+
+def add_sequence_io(extension: str, sequence_io: SequenceIO) -> None:
+  """Adds a record IO system with a prefix."""
+  _registry.add(extension, sequence_io)
+
+
+def _resolve_path(path: Union[str, os.PathLike[str]]) -> str:
+  if isinstance(path, str):
+    return path
+  elif hasattr(path, '__fspath__'):
+    return path.__fspath__()
+  else:
+    raise ValueError(f'Unsupported path: {path!r}.')
+
+
+def open_sequence(
+    path: Union[str, os.PathLike[str]],
+    mode: str = 'r',
+    *,
+    serializer: Optional[
+        Callable[[Any], Union[bytes, str]]
+    ] = None,
+    deserializer: Optional[
+        Callable[[Union[bytes, str]], Any]
+    ] = None,
+) -> Sequence:
+  """Open sequence for reading or writing.
+
+  Args:
+    path: The path to the sequence.
+    mode: The mode of the sequence.
+    serializer: (Optional) A serializer function for converting a structured
+      object to a string or bytes.
+    deserializer: (Optional) A deserializer function for converting a string or
+      bytes to a structured object.
+
+  Returns:
+    A sequence for reading or writing.
+  """
+  return _registry.get(path).open(
+      path, mode, serializer=serializer, deserializer=deserializer
+  )
+
+
+class MemorySequence(Sequence):
+  """An in-memory sequence."""
+
+  def __init__(
+      self,
+      path: str,
+      mode: str,
+      records: list[Union[str, bytes]],
+      serializer: Optional[Callable[[Any], Union[bytes, str]]],
+      deserializer: Optional[Callable[[Union[bytes, str]], Any]]
+  ):
+    super().__init__(serializer, deserializer)
+    self._path = path
+    self._mode = mode
+    self._records = records
+    self._closed = False
+
+  def _add(self, record: Union[str, bytes]) -> None:
+    if 'w' not in self._mode and 'a' not in self._mode:
+      raise ValueError(
+          f'Cannot write record {record!r} to memory sequence {self._path!r} '
+          f'with mode {self._mode!r}.'
+      )
+    if self._closed:
+      raise ValueError(
+          f'Cannot write record {record!r} to a closed writer for '
+          f'{self._path!r}.'
+      )
+    self._records.append(record)
+
+  def __len__(self):
+    return len(self._records)
+
+  def _iter(self):
+    if 'r' not in self._mode:
+      raise ValueError(
+          f'Cannot read memory sequence {self._path!r} with '
+          f'mode {self._mode!r}.'
+      )
+    if self._closed:
+      raise ValueError(
+          f'Cannot iterate over a closed sequence reader {self._path!r}.'
+      )
+    return iter(self._records)
+
+  def flush(self):
+    pass
+
+  def close(self) -> None:
+    self._closed = True
+
+
+class MemorySequenceIO(SequenceIO):
+  """Memory-based record IO."""
+
+  def __init__(self):
+    super().__init__()
+    self._root = collections.defaultdict(list)
+
+  def open(
+      self,
+      path: Union[str, os.PathLike[str]],
+      mode: str,
+      *,
+      serializer: Optional[Callable[[Any], Union[bytes, str]]],
+      deserializer: Optional[Callable[[Union[bytes, str]], Any]],
+      **kwargs
+  ) -> Sequence:
+    """Opens a reader for a sequence."""
+    del kwargs
+    if 'w' in mode:
+      self._root[path] = []
+    return MemorySequence(
+        path, mode, self._root[path],
+        serializer=serializer, deserializer=deserializer
+    )
+
+
+add_sequence_io('mem', MemorySequenceIO())
+
+
+class LineSequence(Sequence):
+  """A new-line broken sequence."""
+
+  def __init__(
+      self,
+      file: file_system.File,
+      serializer: Optional[Callable[[Any], Union[bytes, str]]],
+      deserializer: Optional[Callable[[Union[bytes, str]], Any]],
+  ) -> None:
+    super().__init__(serializer, deserializer)
+    self._file = file
+
+  def __len__(self):
+    raise NotImplementedError(
+        '__len__ is not supported for LineSequence. '
+        'Use `len(list(iter(sequence)))` instead.'
+    )
+
+  def _iter(self):
+    while True:
+      line = self._file.readline()
+      if not line:
+        break
+      yield line.rstrip('\n')
+
+  def _add(self, record: Union[str, bytes]) -> None:
+    self._file.write(record.rstrip('\n'))
+    self._file.write('\n')
+
+  def flush(self) -> None:
+    self._file.flush()
+
+  def close(self) -> None:
+    self._file.close()
+
+
+class LineSequenceIO(SequenceIO):
+  """Line-based record IO."""
+
+  def open(
+      self,
+      path: Union[str, os.PathLike[str]],
+      mode: str,
+      *,
+      serializer: Optional[Callable[[Any], Union[bytes, str]]],
+      deserializer: Optional[Callable[[Union[bytes, str]], Any]],
+      **kwargs
+  ) -> Sequence:
+    """Opens a reader for a sequence."""
+    del kwargs
+    return LineSequence(
+        file_system.open(path, mode), serializer, deserializer
+    )
+

pyglove/core/io/sequence_test.py

@@ -0,0 +1,137 @@
+# Copyright 2024 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 os
+import pathlib
+import tempfile
+import unittest
+from pyglove.core.io import sequence as sequence_io
+# We need to import this module to register the default parser/serializer.
+import pyglove.core.symbolic as pg_symbolic
+
+
+class LineSequenceIOTest(unittest.TestCase):
+
+  def test_read_write(self):
+    tmp_dir = tempfile.gettempdir()
+    file1 = os.path.join(tmp_dir, 'file1')
+    with pg_symbolic.open_jsonl(file1, 'w') as f:
+      self.assertIsInstance(f, sequence_io.LineSequence)
+      f.add(1)
+      f.add(' foo')
+      f.add(' bar ')
+      f.flush()
+      self.assertTrue(os.path.exists(file1))
+
+      f.add('baz\n')
+      f.add(dict(x=1))
+
+    self.assertTrue(os.path.exists(file1))
+    with pg_symbolic.open_jsonl(file1, 'r') as f:
+      self.assertIsInstance(f, sequence_io.LineSequence)
+      self.assertEqual(list(iter(f)), [1, ' foo', ' bar ', 'baz\n', dict(x=1)])
+
+    with pg_symbolic.open_jsonl(file1, 'a') as f:
+      f.add('qux')
+
+    with pg_symbolic.open_jsonl(file1, 'r') as f:
+      with self.assertRaisesRegex(
+          NotImplementedError, '__len__ is not supported'
+      ):
+        _ = len(f)
+      self.assertEqual(
+          list(iter(f)), [1, ' foo', ' bar ', 'baz\n', dict(x=1), 'qux']
+      )
+
+  def test_read_write_with_raw_texts(self):
+    tmp_dir = tempfile.gettempdir()
+    file2 = os.path.join(tmp_dir, 'file2')
+    with sequence_io.open_sequence(file2, 'w') as f:
+      self.assertIsInstance(f, sequence_io.LineSequence)
+      with self.assertRaisesRegex(
+          ValueError, 'Cannot write record with type'
+      ):
+        f.add(1)
+      f.add('foo\nbar\n')
+
+    with sequence_io.open_sequence(file2, 'r') as f:
+      with self.assertRaisesRegex(
+          ValueError, 'not writable'
+      ):
+        f.add('baz')
+      self.assertEqual(list(iter(f)), ['foo', 'bar'])
+
+
+class MemoryRecordIOIOTest(unittest.TestCase):
+
+  def test_resolve_path(self):
+    self.assertEqual(
+        sequence_io._resolve_path('/file1.mem@123'),
+        '/file1.mem@123'
+    )
+    self.assertEqual(
+        sequence_io._resolve_path(pathlib.Path('/file1.mem@123.txt')),
+        '/file1.mem@123.txt'
+    )
+    with self.assertRaisesRegex(ValueError, 'Unsupported path'):
+      sequence_io._resolve_path(1)
+
+  def test_read_write(self):
+    with sequence_io.open_sequence('/file1.mem@123', 'w') as f:
+      self.assertIsInstance(f, sequence_io.MemorySequence)
+      f.add(' foo')
+      f.add(' bar ')
+      f.flush()
+      f.add('baz')
+      with self.assertRaisesRegex(
+          ValueError, 'Cannot write record with type'
+      ):
+        f.add(1)
+      with self.assertRaisesRegex(
+          ValueError, 'Cannot read memory sequence'
+      ):
+        next(iter(f))
+
+    with self.assertRaisesRegex(
+        ValueError, 'Cannot write record .* to a closed writer'
+    ):
+      f.add('qux')
+
+    with sequence_io.open_sequence('/file1.mem@123', 'a') as f:
+      self.assertIsInstance(f, sequence_io.MemorySequence)
+      f.add('qux')
+
+    with sequence_io.open_sequence('/file1.mem@123') as f:
+      self.assertIsInstance(f, sequence_io.MemorySequence)
+      self.assertEqual(len(f), 4)
+      self.assertEqual(list(f), [' foo', ' bar ', 'baz', 'qux'])
+      with self.assertRaisesRegex(
+          ValueError, 'Cannot write record .* to memory sequence'
+      ):
+        f.add('abc')
+
+    with self.assertRaisesRegex(
+        ValueError, 'Cannot iterate over a closed sequence reader'
+    ):
+      next(iter(f))
+
+    with sequence_io.open_sequence('/file1.mem@123', 'w') as f:
+      f.add('abc')
+
+    with sequence_io.open_sequence('/file1.mem@123', 'r') as f:
+      self.assertEqual(list(iter(f)), ['abc'])
+
+
+if __name__ == '__main__':
+  unittest.main()

pyglove/core/object_utils/__init__.py

@@ -154,10 +154,11 @@ from pyglove.core.object_utils.docstr_utils import docstr
 # Handling exceptions.
 from pyglove.core.object_utils.error_utils import catch_errors
 from pyglove.core.object_utils.error_utils import CatchErrorsContext
+from pyglove.core.object_utils.error_utils import ErrorInfo
 
-# Profiling
-from pyglove.core.object_utils.profiling import timeit
-from pyglove.core.object_utils.profiling import TimeIt
+# Timing.
+from pyglove.core.object_utils.timing import timeit
+from pyglove.core.object_utils.timing import TimeIt
 
 # pylint: enable=g-importing-member
 # pylint: enable=g-bad-import-order

pyglove/core/object_utils/error_utils.py

@@ -17,7 +17,71 @@ 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.object_utils import formatting
+from pyglove.core.object_utils import json_conversion
+
+
+@dataclasses.dataclass(frozen=True)
+class ErrorInfo(json_conversion.JSONConvertible, formatting.Formattable):
+  """Serializable error information.
+
+  Attributes:
+    type: The type path of the error. 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.
+  """
+
+  type: str
+  description: str
+  stacktrace: str
+
+  @classmethod
+  def _get_type(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(
+        type=cls._get_type(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(
+            type=(self.type, None),
+            description=(self.description, None),
+            stacktrace=(self.stacktrace, None),
+        ),
+        exclude_default=True,
+        **kwargs,
+    )
+
+  def format(self, *args, **kwargs) -> str:
+    return formatting.kvlist_str(
+        [
+            ('type', self.type, None),
+            ('description', self.description, None),
+            ('stacktrace', self.stacktrace, None),
+        ],
+        *args,
+        label=self.__class__.__name__,
+        **kwargs,
+    )
 
 
 @dataclasses.dataclass()

pyglove/core/object_utils/error_utils_test.py

@@ -11,12 +11,66 @@
 # 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
 
 
+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.type, '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(
+        type='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(
+        type='ValueError.ZeroDivisionError',
+        description='Bad call to `foo`',
+        stacktrace='Traceback (most recent call last)',
+    )
+    self.assertEqual(
+        error_info.format(compact=False),
+        inspect.cleandoc(
+            """
+            ErrorInfo(
+              type='ValueError.ZeroDivisionError',
+              description='Bad call to `foo`',
+              stacktrace='Traceback (most recent call last)'
+            )
+            """
+        )
+    )
+
+
 class CatchErrorsTest(unittest.TestCase):
 
   def assert_caught_error(self, func, errors_to_catch):

pyglove/core/object_utils/{profiling.py → timing.py}

@@ -11,12 +11,13 @@
 # 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.
-"""Utilities for profiling."""
+"""Utilities for timing."""
 
 import dataclasses
 import time
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
+from pyglove.core.object_utils import json_conversion
 from pyglove.core.object_utils import thread_local
 
 
@@ -24,12 +25,12 @@ class TimeIt:
   """Context manager for timing the execution of a code block."""
 
   @dataclasses.dataclass(frozen=True)
-  class Status:
+  class Status(json_conversion.JSONConvertible):
     """Status of a single `pg.timeit`."""
     name: str
-    elapse: float
-    has_ended: bool
-    error: Optional[Exception]
+    elapse: float = 0.0
+    has_ended: bool = True
+    error: Optional[Exception] = None
 
     @property
     def has_started(self) -> bool:
@@ -41,12 +42,24 @@ class TimeIt:
       """Returns whether the context has error."""
       return self.error is not None
 
+    def to_json(self, **kwargs) -> Dict[str, Any]:
+      return self.to_json_dict(
+          fields=dict(
+              name=(self.name, None),
+              elapse=(self.elapse, 0.0),
+              has_ended=(self.has_ended, True),
+              error=(self.error, None),
+          ),
+          exclude_default=True,
+          **kwargs,
+      )
+
   @dataclasses.dataclass
-  class StatusSummary:
+  class StatusSummary(json_conversion.JSONConvertible):
     """Aggregated summary for repeated calls for `pg.timeit`."""
 
     @dataclasses.dataclass
-    class Entry:
+    class Entry(json_conversion.JSONConvertible):
       """Aggregated status from the `pg.timeit` calls of the same name."""
 
       num_started: int = 0
@@ -65,16 +78,41 @@ class TimeIt:
         if status.has_error:
           self.num_failed += 1
 
+      def to_json(self, **kwargs) -> Dict[str, Any]:
+        return self.to_json_dict(
+            fields=dict(
+                num_started=(self.num_started, 0),
+                num_ended=(self.num_ended, 0),
+                num_failed=(self.num_failed, 0),
+                avg_duration=(self.avg_duration, 0.0),
+            ),
+            exclude_default=True,
+            **kwargs,
+        )
+
     breakdown: dict[str, 'TimeIt.StatusSummary.Entry'] = (
         dataclasses.field(default_factory=dict)
     )
 
-    def aggregate(self, timeit_obj: 'TimeIt'):
-      for k, v in timeit_obj.status().items():
+    def __bool__(self) -> bool:
+      """Returns True if the summary is non-empty."""
+      return bool(self.breakdown)
+
+    def aggregate(self, timeit_status: Dict[str, 'TimeIt.Status']):
+      for k, v in timeit_status.items():
         if k not in self.breakdown:
           self.breakdown[k] = TimeIt.StatusSummary.Entry()
         self.breakdown[k].update(v)
 
+    def to_json(self, **kwargs) -> Dict[str, Any]:
+      return self.to_json_dict(
+          fields=dict(
+              breakdown=(self.breakdown, {}),
+          ),
+          exclude_default=True,
+          **kwargs,
+      )
+
   def __init__(self, name: str):
     self._name = name
     self._start_time = None

pyglove/core/object_utils/{profiling_test.py → timing_test.py}

@@ -15,14 +15,15 @@
 import time
 import unittest
 
-from pyglove.core.object_utils import profiling
+from pyglove.core.object_utils import json_conversion
+from pyglove.core.object_utils import timing
 
 
 class TimeItTest(unittest.TestCase):
   """Tests for `pg.symbolic.thread_local`."""
 
   def test_basics(self):
-    tc = profiling.TimeIt('node')
+    tc = timing.TimeIt('node')
     self.assertFalse(tc.has_started)
     self.assertEqual(tc.elapse, 0)
 
@@ -37,7 +38,7 @@ class TimeItTest(unittest.TestCase):
     self.assertEqual(tc.elapse, elapse1)
 
   def test_timeit(self):
-    with profiling.timeit('node') as t:
+    with timing.timeit('node') as t:
       self.assertEqual(t.name, 'node')
       self.assertIsNotNone(t.start_time)
       self.assertTrue(t.has_started)
@@ -46,9 +47,9 @@ class TimeItTest(unittest.TestCase):
       time.sleep(0.5)
       elapse1 = t.elapse
       self.assertEqual(t.children, [])
-      with profiling.timeit('child') as t1:
+      with timing.timeit('child') as t1:
         time.sleep(0.5)
-        with profiling.timeit('grandchild') as t2:
+        with timing.timeit('grandchild') as t2:
           time.sleep(0.5)
 
         self.assertEqual(t.children, [t1])
@@ -63,7 +64,7 @@ class TimeItTest(unittest.TestCase):
         self.assertTrue(r['node.child.grandchild'].has_ended)
 
         with self.assertRaisesRegex(ValueError, '.* already exists'):
-          with profiling.timeit('grandchild'):
+          with timing.timeit('grandchild'):
             pass
 
     elapse2 = t.elapse
@@ -83,12 +84,17 @@ class TimeItTest(unittest.TestCase):
     )
     self.assertTrue(all(v.has_ended for v in statuss.values()))
     self.assertFalse(any(v.has_error for v in statuss.values()))
+    status = t.status()
+    json_dict = json_conversion.to_json(status)
+    status2 = json_conversion.from_json(json_dict)
+    self.assertIsNot(status2, status)
+    self.assertEqual(status2, status)
 
   def test_timeit_with_error(self):
     with self.assertRaises(ValueError):
-      with profiling.timeit('node') as t:
-        with profiling.timeit('child') as t1:
-          with profiling.timeit('grandchild') as t2:
+      with timing.timeit('node') as t:
+        with timing.timeit('child') as t1:
+          with timing.timeit('grandchild') as t2:
             raise ValueError('error')
 
     r = t.status()
@@ -102,20 +108,22 @@ class TimeItTest(unittest.TestCase):
     self.assertTrue(t2.has_error)
 
   def test_timeit_summary(self):
-    summary = profiling.TimeIt.StatusSummary()
+    summary = timing.TimeIt.StatusSummary()
+    self.assertFalse(summary)
     for i in range(10):
-      with profiling.timeit('node') as t:
+      with timing.timeit('node') as t:
         time.sleep(0.1)
-        with profiling.timeit('child'):
+        with timing.timeit('child'):
           time.sleep(0.1)
           try:
-            with profiling.timeit('grandchild'):
+            with timing.timeit('grandchild'):
               time.sleep(0.1)
               if i < 2:
                 raise ValueError('error')
           except ValueError:
             pass
-      summary.aggregate(t)
+      summary.aggregate(t.status())
+    self.assertTrue(summary)
     self.assertEqual(
         list(summary.breakdown.keys()),
         ['node', 'node.child', 'node.child.grandchild']
@@ -132,6 +140,11 @@ class TimeItTest(unittest.TestCase):
         [x.num_failed for x in summary.breakdown.values()],
         [0, 0, 2]
     )
+    # Test serialization.
+    json_dict = summary.to_json()
+    summary2 = timing.TimeIt.StatusSummary.from_json(json_dict)
+    self.assertIsNot(summary2, summary)
+    self.assertEqual(summary2.breakdown, summary.breakdown)
 
 
 if __name__ == '__main__':

pyglove/core/symbolic/__init__.py

@@ -120,6 +120,7 @@ from pyglove.core.symbolic.base import to_json
 from pyglove.core.symbolic.base import to_json_str
 from pyglove.core.symbolic.base import load
 from pyglove.core.symbolic.base import save
+from pyglove.core.symbolic.base import open_jsonl
 
 # Interfaces for pure symbolic objects.
 from pyglove.core.symbolic.pure_symbolic import PureSymbolic

pyglove/core/symbolic/base.py

@@ -2310,6 +2310,42 @@ def save(value: Any, path: str, *args, **kwargs) -> Any:
   return save_handler(value, path, *args, **kwargs)
 
 
+def open_jsonl(
+    path: str,
+    mode: str = 'r',
+    **kwargs
+) -> pg_io.Sequence:
+  """Open a JSONL file for reading or writing.
+
+  Example::
+
+    with pg.open_jsonl('my_file.jsonl', 'w') as f:
+      f.add(1)
+      f.add('foo')
+      f.add(dict(x=1))
+
+    with pg.open_jsonl('my_file.jsonl', 'r') as f:
+      for value in f:
+        print(value)
+
+  Args:
+    path: The path to the file.
+    mode: The mode of the file.
+    **kwargs: Additional keyword arguments that will be passed to
+      ``pg_io.open_sequence``.
+
+  Returns:
+    A sequence for PyGlove objects.
+  """
+  return pg_io.open_sequence(
+      path,
+      mode,
+      serializer=to_json_str,
+      deserializer=from_json_str,
+      **kwargs
+  )
+
+
 def default_load_handler(
     path: str,
     file_format: Literal['json', 'txt'] = 'json',

pyglove/core/views/html/base.py

@@ -333,7 +333,11 @@ class Html(base.Content):
     return s
 
   @classmethod
-  def escape(cls, s: WritableTypes) -> WritableTypes:
+  def escape(
+      cls,
+      s: WritableTypes,
+      javascript_str: bool = False
+  ) -> WritableTypes:
     """Escapes an HTML writable object."""
     if s is None:
       return None
@@ -341,11 +345,22 @@ class Html(base.Content):
     if callable(s):
       s = s()
 
-    if isinstance(s, str):
+    def _escape(s: str) -> str:
+      if javascript_str:
+        return (
+            s.replace('\\', '\\\\')
+            .replace('"', '\\"')
+            .replace('\r', '\\r')
+            .replace('\n', '\\n')
+            .replace('\t', '\\t')
+        )
       return html_lib.escape(s)
+
+    if isinstance(s, str):
+      return _escape(s)
     else:
       assert isinstance(s, Html), s
-      return Html(html_lib.escape(s.content)).write(
+      return Html(_escape(s.content)).write(
           s, shared_parts_only=True
       )
 

pyglove/core/views/html/base_test.py

@@ -695,6 +695,8 @@ class HtmlTest(TestCase):
     self.assertEqual(Html.escape('foo"bar'), 'foo"bar')
     self.assertEqual(Html.escape(Html('foo"bar')), Html('foo"bar'))
     self.assertEqual(Html.escape(lambda: 'foo"bar'), 'foo"bar')
+    self.assertEqual(Html.escape('"x=y"', javascript_str=True), '\\"x=y\\"')
+    self.assertEqual(Html.escape('x\n"', javascript_str=True), 'x\\n\\"')
 
   def test_concate(self):
     self.assertIsNone(Html.concate(None))

pyglove/core/views/html/tree_view.py

@@ -1215,6 +1215,7 @@ class HtmlTreeView(HtmlView):
       parent: Any = None,
       root_path: Optional[KeyPath] = None,
       css_classes: Optional[Sequence[str]] = None,
+      id: Optional[str] = None,  # pylint: disable=redefined-builtin
       content: Union[str, Html, None] = None,
       **kwargs,
   ) -> Html:
@@ -1225,8 +1226,9 @@ class HtmlTreeView(HtmlView):
       parent: The parent of the value.
       root_path: The root path of the value.
       css_classes: CSS classes to add to the HTML element.
+      id: The ID of the tooltip span element. If None, no ID will be added.
       content: The content to render. If None, the value will be rendered.
-      **kwargs: Additional keyword arguments passed from the user that 
+      **kwargs: Additional keyword arguments passed from the user that
         will be ignored.
 
     Returns:
@@ -1248,6 +1250,7 @@ class HtmlTreeView(HtmlView):
     return Html.element(
         'span',
         [content],
+        id=id,
         css_classes=[
             'tooltip',
             css_classes,
@@ -1266,6 +1269,9 @@ class HtmlTreeView(HtmlView):
           position: absolute;
           z-index: 1;
         }
+        span.tooltip:hover {
+          visibility: visible;
+        }
         """
     )
 

pyglove/core/views/html/tree_view_test.py

@@ -74,6 +74,9 @@ class TooltipTest(TestCase):
           position: absolute;
           z-index: 1;
         }
+        span.tooltip:hover {
+          visibility: visible;
+        }
         </style>
         """
     )
@@ -122,6 +125,9 @@ class SummaryTest(TestCase):
           position: absolute;
           z-index: 1;
         }
+        span.tooltip:hover {
+          visibility: visible;
+        }
         /* Summary styles. */
         details.pyglove summary {
           font-weight: bold;
@@ -314,6 +320,9 @@ class ContentTest(TestCase):
           position: absolute;
           z-index: 1;
         }
+        span.tooltip:hover {
+          visibility: visible;
+        }
         /* Summary styles. */
         details.pyglove summary {
           font-weight: bold;

{pyglove-0.4.5.dev202410290809.dist-info → pyglove-0.4.5.dev202410301815.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pyglove
-Version: 0.4.5.dev202410290809
+Version: 0.4.5.dev202410301815
 Summary: PyGlove: A library for manipulating Python objects.
 Home-page: https://github.com/google/pyglove
 Author: PyGlove Authors

{pyglove-0.4.5.dev202410290809.dist-info → pyglove-0.4.5.dev202410301815.dist-info}/RECORD

@@ -1,5 +1,5 @@
 pyglove/__init__.py,sha256=YPUWALRDu4QuI7N3TUmVaVr513hUt0qXrjY582I1q5c,1352
-pyglove/core/__init__.py,sha256=tHz9DcjN_vGIaYPt4TDWiJ9XXfQmKXsMR7tYwQaoUX4,8959
+pyglove/core/__init__.py,sha256=omnsY79mvCd39G9HKtuYHltlgwB5kTdgqKQAoXWgT-I,8992
 pyglove/core/logging.py,sha256=Sr5nLyUQmc4CAEYAq8qbP3ghpjmpz2sOuhq2A0tgQ6I,2456
 pyglove/core/logging_test.py,sha256=QM59tCAFOUKm4gM6ef_9eG1AvSGEYRDLrrprUQyslcA,1956
 pyglove/core/detouring/__init__.py,sha256=ck_n2VSuU31HNVYQkbG4Zvnx90mNYtSVc2StN3rXbU8,1891
@@ -42,18 +42,20 @@ pyglove/core/hyper/numerical.py,sha256=ifOdw4G9xukeELEiB8__zsR0X6O06_k9isq0i6YZw
 pyglove/core/hyper/numerical_test.py,sha256=_xZe0qe3wJdrQL4F8pgt0zf04zpbc_9c_5va3mdd5_8,4259
 pyglove/core/hyper/object_template.py,sha256=x6aBe-6ueo3CsFJ2yhbXbcdJKfId6wm_KsgkHHbpSoY,22325
 pyglove/core/hyper/object_template_test.py,sha256=TEFX7LIqUvdCdJILnK_gP5xIgNJKzRnioUF0CGVBzcY,9105
-pyglove/core/io/__init__.py,sha256=_wqJqPGny8LcVOze2ym_5q8S5OGTZRVpTsQ0_DAmJvY,746
-pyglove/core/io/file_system.py,sha256=q-XDwk9A2oWX90OaZf7CedpNwu_aPPKbHQC2mCwNo4A,13084
-pyglove/core/io/file_system_test.py,sha256=I4tlj7oi8jWXEYVfPGimhHYL9s_nHVrxZsSbwOyeFSc,8659
-pyglove/core/object_utils/__init__.py,sha256=AghdwjgQQIB9lu90aWeBUYHP3CrrolX1YiPgXNMeVfg,8706
+pyglove/core/io/__init__.py,sha256=4ZT1a595DqQuLTNYc2JP_eCp_KesXvHmKRkr777bzpg,785
+pyglove/core/io/file_system.py,sha256=KR1riWNBVZpkgIcqPUhR5xqjCN5Lhp937x7Ubjpv-jE,13283
+pyglove/core/io/file_system_test.py,sha256=FX0ySuh_Xcg1RO68do7ikD4pvslKUzfSpwZ6P4wIP7c,8691
+pyglove/core/io/sequence.py,sha256=Kir8JM8lxjrlKr1gs3hvhCB8F08PPz4K8-aDI_US9Vo,8042
+pyglove/core/io/sequence_test.py,sha256=iqDR-MPm64VbUPKFYocVKH91Y6Z6tP1TZB0PzCn92tE,4265
+pyglove/core/object_utils/__init__.py,sha256=bFyt-_uNzpOjJGDfYFTRP1kV6rL4du9xGx_OaFMcIxI,8758
 pyglove/core/object_utils/codegen.py,sha256=HQlpjz42oh_2lrGSflwuAcEneE9Bv8pgPahQzvJm1KI,1568
 pyglove/core/object_utils/codegen_test.py,sha256=ThPVyE805plNvM8vwlDxMrcjV-GGdDgbBLSqSOpCvJ0,2166
 pyglove/core/object_utils/common_traits.py,sha256=9wfWvcq3uuVgL8mQexgJmWdF2aaQ-wWKzLXO9dqaqLk,3621
 pyglove/core/object_utils/common_traits_test.py,sha256=wIeVsF3kon9K0Kbblcaib9hBJeZ76LyGZDD-5A1BD9w,1182
 pyglove/core/object_utils/docstr_utils.py,sha256=5BY40kXozPKVGOB0eN8jy1P5_GHIzqFJ9FXAu_kzxaw,5119
 pyglove/core/object_utils/docstr_utils_test.py,sha256=1NIsXXpp87HGC8LB7vx1KEMVVrmeGi733Xpuf6yTUq0,4287
-pyglove/core/object_utils/error_utils.py,sha256=4uZxi9dGYVsJgvctt8dFnbpbfT8Kqma6KwqjHJ5okn8,3949
-pyglove/core/object_utils/error_utils_test.py,sha256=YOL2gL5_v_KspA49ulyYV3Vse5jIFfYHYxoq09GGKdg,2605
+pyglove/core/object_utils/error_utils.py,sha256=wnGX_aR4TdbCylvUOh-x6PPYzGSJBnfyX9MANpk03u8,5780
+pyglove/core/object_utils/error_utils_test.py,sha256=vlFkR41jNLzUnwqF6IXy7HmXhpLlwOvUvlVMoMSOitk,4175
 pyglove/core/object_utils/formatting.py,sha256=cXGWJd6zleS9QiEfY3w4fqWfFYgi_QF8S2IapTOzAj8,14994
 pyglove/core/object_utils/formatting_test.py,sha256=QGbpUS8u70NoF8-3v5IYmPWEsML3ZbIwkQpUMwQdTkQ,13589
 pyglove/core/object_utils/hierarchical.py,sha256=za9sA6XnlUOKUHMKNwQPNSOK9E3wDqXFQB0AND2Yghw,19754
@@ -62,10 +64,10 @@ pyglove/core/object_utils/json_conversion.py,sha256=1XEwz4LzEh8wHsTIC7jL7lD7L5tK
 pyglove/core/object_utils/json_conversion_test.py,sha256=KGt0r628KHi4fTfeCgrhirfINpYXIR45l5ER4z09jjI,11907
 pyglove/core/object_utils/missing.py,sha256=0liMs9iEyQxxu6UohdJ5hEM246e9Nu05qp0hqriHsl0,1459
 pyglove/core/object_utils/missing_test.py,sha256=B36p-vqUvAnXWMszAj9GOPBN0_8cq7vVF61AkcsZ9qU,1396
-pyglove/core/object_utils/profiling.py,sha256=hNMd4BY2jCOOV3o0VTlweHxZwd8yGylDzW6GgHU1-j8,5385
-pyglove/core/object_utils/profiling_test.py,sha256=LvFy8qrAJUD-A4O7NBWZJ5YfjSGtc_zpmATCMs9hSp8,4322
 pyglove/core/object_utils/thread_local.py,sha256=i-CnyY3VREtLfAj4_JndBnsKuQLIgwG29ma8dAyRxbI,4839
 pyglove/core/object_utils/thread_local_test.py,sha256=EvU1-TF7KqpLQxxBvHd7dxtuY22YUQSIwQ0UcR-NORA,6816
+pyglove/core/object_utils/timing.py,sha256=PeSN59g8vQDpenPgsWf4oYP45c0UDSpnFYDKhrtcSBc,6665
+pyglove/core/object_utils/timing_test.py,sha256=Ye4qBrNl5ix2C_K-qT5TxKHui4Kvwmea3uUTx8t26qk,4828
 pyglove/core/object_utils/value_location.py,sha256=lSFQNTazY2M6_nRLmMbouqZAcZSiOLZQnmQPMD2FDMs,26770
 pyglove/core/object_utils/value_location_test.py,sha256=wCZYqWf_tq3l4ZM6dnkAs1obmq8LfYawkns2NKos9kk,21406
 pyglove/core/patching/__init__.py,sha256=C1Q1cWPV74YL3eXbzGvc-8aPw1DR8EK6lRhQYDCwHek,2059
@@ -75,8 +77,8 @@ pyglove/core/patching/pattern_based.py,sha256=g56J19s6uTsaX5mHt--TnGifUQCjYjWpLX
 pyglove/core/patching/pattern_based_test.py,sha256=PW1EcVfsFPB6wtgwg3s4dzvigWn3b5S8eMNGo0SJiZ0,2771
 pyglove/core/patching/rule_based.py,sha256=QJsDsYVpyfUoaTZeXZtAdFuxpvB1eniY0qYdZvVqzk0,17056
 pyglove/core/patching/rule_based_test.py,sha256=qfy0ILmczV_LMHWEnwo2y079OrJsGYO0nKxSZdmIUcI,18782
-pyglove/core/symbolic/__init__.py,sha256=p9NvBRakrFQl_jGQFglFwOblxu9B3OOPs34xTLCscMw,5737
-pyglove/core/symbolic/base.py,sha256=WMdwqVriVOwuPBXp8qkXh69nenvNPs__0FWRIT5cyEo,77579
+pyglove/core/symbolic/__init__.py,sha256=jYq-LwR1Ql3iMChjz9lN-0heDKiocmxR0ZJFEJ8RHHQ,5787
+pyglove/core/symbolic/base.py,sha256=JNgh6eU_d5DRKQKlw4QP0FSJJg1RaVTrgGBjumg3ELs,78292
 pyglove/core/symbolic/base_test.py,sha256=YLMlQGF4TbxAB2PcN1-ckI70gRIVDxabv-goh43CV7A,7333
 pyglove/core/symbolic/boilerplate.py,sha256=YO8ZTZJ3VfAqeHqJpIY_ORVDxc1XIzxdplNlpexWEyc,6000
 pyglove/core/symbolic/boilerplate_test.py,sha256=1CZ1W6kq3l-3tpaknhGFa04V18bO7vPzis5qzWnxHEs,5252
@@ -141,10 +143,10 @@ pyglove/core/views/__init__.py,sha256=hPqM3agvS3skQmLmZyV-5ai9vkhzJI65RLXmS9P-mO
 pyglove/core/views/base.py,sha256=yuDHX5g1MgqyeVJLJPpWeh33opEafPygV1fjIOshWLE,26163
 pyglove/core/views/base_test.py,sha256=F6nou7reS_Kmr2H57MUgcnc1bIp9Z3BWCvHOvMqjGkQ,16642
 pyglove/core/views/html/__init__.py,sha256=FpwzhpomAkIjaOQh6mkIjaXC0KWHvGFIRCAQUAdxdC8,1029
-pyglove/core/views/html/base.py,sha256=i87aldEsw8V6vW2RkmVyCPX70R_yh0-gwD45gbW1e-4,15513
-pyglove/core/views/html/base_test.py,sha256=JtzM22QPLINW9YvOuSCtj_eXsDWD9zG8VVf_LJHgfxw,22192
-pyglove/core/views/html/tree_view.py,sha256=orhMI7zOFKBY5v0Bhca700ulKsVKXI_Rj90kzpiz56U,51245
-pyglove/core/views/html/tree_view_test.py,sha256=7T0kjlEW2DxUKsfOnFU_9xmgP6DM4Rq2Mp32teHNYZs,74273
+pyglove/core/views/html/base.py,sha256=ob_wyUhlrZNQfrzeF5BWr8Bt8PH7zkYBwvR0xa40ABQ,15837
+pyglove/core/views/html/base_test.py,sha256=Oko2B0vw3NYzJplEK4rhFoizOPeuBuebYZU-k39tMX0,22343
+pyglove/core/views/html/tree_view.py,sha256=MM3Dct8kujaubdK_99KMwKJx7GFjLxsyTJV0sewH20M,51474
+pyglove/core/views/html/tree_view_test.py,sha256=AzStpY_iuvnJUWViuFXFGyrGbNv3P2qVFQyw105O1Jg,74483
 pyglove/ext/__init__.py,sha256=3jp8cJvKW6PENOZlmVAbT0w-GBRn_kjhc0wDX3XjpOE,755
 pyglove/ext/early_stopping/__init__.py,sha256=_xkT3K_MycFuF3k2N74sfr9FPFcXMD3pfv5vuc4Tx5M,1116
 pyglove/ext/early_stopping/base.py,sha256=KolexwKDAXhI9iY7hoaDO1mbc3TA4gDn0UifxURwbk0,2358
@@ -184,8 +186,8 @@ pyglove/ext/scalars/randoms.py,sha256=LkMIIx7lOq_lvJvVS3BrgWGuWl7Pi91-lA-O8x_gZs
 pyglove/ext/scalars/randoms_test.py,sha256=nEhiqarg8l_5EOucp59CYrpO2uKxS1pe0hmBdZUzRNM,2000
 pyglove/ext/scalars/step_wise.py,sha256=IDw3tuTpv0KVh7AN44W43zqm1-E0HWPUlytWOQC9w3Y,3789
 pyglove/ext/scalars/step_wise_test.py,sha256=TL1vJ19xVx2t5HKuyIzGoogF7N3Rm8YhLE6JF7i0iy8,2540
-pyglove-0.4.5.dev202410290809.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-pyglove-0.4.5.dev202410290809.dist-info/METADATA,sha256=IioX481iPq1LaxeDpORWZVVgaV3mxy3WKlpGM8ibHb0,6666
-pyglove-0.4.5.dev202410290809.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
-pyglove-0.4.5.dev202410290809.dist-info/top_level.txt,sha256=wITzJSKcj8GZUkbq-MvUQnFadkiuAv_qv5qQMw0fIow,8
-pyglove-0.4.5.dev202410290809.dist-info/RECORD,,
+pyglove-0.4.5.dev202410301815.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+pyglove-0.4.5.dev202410301815.dist-info/METADATA,sha256=jIeDOiGyeIig84equ3ZG2EjO4s1his-CfD1PjMETKpU,6666
+pyglove-0.4.5.dev202410301815.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+pyglove-0.4.5.dev202410301815.dist-info/top_level.txt,sha256=wITzJSKcj8GZUkbq-MvUQnFadkiuAv_qv5qQMw0fIow,8
+pyglove-0.4.5.dev202410301815.dist-info/RECORD,,

{pyglove-0.4.5.dev202410290809.dist-info → pyglove-0.4.5.dev202410301815.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.2.0)
+Generator: setuptools (75.3.0)
 Root-Is-Purelib: true
 Tag: py3-none-any