pyglove 0.4.5.dev202410290809__py3-none-any.whl → 0.4.5.dev202411010809__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."""
@@ -112,7 +116,7 @@ class FileSystem(metaclass=abc.ABCMeta):
     """Removes a directory chain based on a path."""
 
 
-def _resolve_path(path: Union[str, os.PathLike[str]]) -> str:
+def resolve_path(path: Union[str, os.PathLike[str]]) -> str:
   if isinstance(path, str):
     return path
   elif hasattr(path, '__fspath__'):
@@ -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)
 
@@ -233,7 +243,7 @@ class MemoryFileSystem(FileSystem):
     self._prefix = prefix
 
   def _internal_path(self, path: Union[str, os.PathLike[str]]) -> str:
-    return '/' + _resolve_path(path).lstrip(self._prefix)
+    return '/' + resolve_path(path).lstrip(self._prefix)
 
   def _locate(self, path: Union[str, os.PathLike[str]]) -> Any:
     current = self._root
@@ -277,7 +287,7 @@ class MemoryFileSystem(FileSystem):
   def _parent_and_name(
       self, path: Union[str, os.PathLike[str]]
   ) -> tuple[dict[str, Any], str]:
-    path = _resolve_path(path)
+    path = resolve_path(path)
     rpos = path.rfind('/')
     assert rpos >= 0, path
     name = path[rpos + 1:]
@@ -372,7 +382,7 @@ class _FileSystemRegistry:
 
   def get(self, path: Union[str, os.PathLike[str]]) -> FileSystem:
     """Gets the file system for a path."""
-    path = _resolve_path(path)
+    path = resolve_path(path)
     for prefix, fs in self._filesystems:
       if path.startswith(prefix):
         return fs

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,299 @@
+# 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 = file_system.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 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,
+    make_dirs_if_not_exist: bool = True,
+) -> 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.
+    make_dirs_if_not_exist: (Optional) Whether to create the directories
+      if they do not exist. Applicable when opening in write or append mode.
+
+  Returns:
+    A sequence for reading or writing.
+  """
+  if 'w' in mode or 'a' in mode:
+    parent_dir = os.path.dirname(path)
+    if make_dirs_if_not_exist:
+      file_system.mkdirs(parent_dir, exist_ok=True)
+  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,124 @@
+# 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 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, 'abc', '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 MemorySequenceIOTest(unittest.TestCase):
+
+  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:
+    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()

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.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):
 
   def assert_caught_error(self, func, errors_to_catch):

pyglove/core/object_utils/json_conversion.py

@@ -202,9 +202,8 @@ class JSONConvertible(metaclass=abc.ABCMeta):
     Returns:
       An instance of cls.
     """
-    del kwargs
     assert isinstance(json_value, dict)
-    init_args = {k: from_json(v) for k, v in json_value.items()
+    init_args = {k: from_json(v, **kwargs) for k, v in json_value.items()
                  if k != JSONConvertible.TYPE_NAME_KEY}
     return cls(**init_args)