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

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/logging_test.py

@@ -11,8 +11,6 @@
 # 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."""
-
 import io
 import logging
 import unittest

pyglove/core/patching/object_factory.py

@@ -14,8 +14,8 @@
 """Object factory based on patchers."""
 
 from typing import Any, Callable, Dict, Optional, Type, Union
-from pyglove.core import object_utils
 from pyglove.core import symbolic
+from pyglove.core import utils
 from pyglove.core.patching import rule_based
 
 
@@ -88,7 +88,7 @@ def ObjectFactory(    # pylint: disable=invalid-name
   # Step 3: Patch with additional parameter override dict if available.
   if params_override:
     value = value.rebind(
-        object_utils.flatten(from_maybe_serialized(params_override, dict)),
-        raise_on_no_change=False)
+        utils.flatten(from_maybe_serialized(params_override, dict)),
+        raise_on_no_change=False,
+    )
   return value
-

pyglove/core/patching/pattern_based.py

@@ -15,8 +15,8 @@
 
 import re
 from typing import Any, Callable, Optional, Tuple, Type, Union
-from pyglove.core import object_utils
 from pyglove.core import symbolic
+from pyglove.core import utils
 
 
 def patch_on_key(
@@ -214,11 +214,11 @@ def patch_on_member(
 
 def _conditional_patch(
     src: symbolic.Symbolic,
-    condition: Callable[
-        [object_utils.KeyPath, Any, symbolic.Symbolic], bool],
+    condition: Callable[[utils.KeyPath, Any, symbolic.Symbolic], bool],
     value: Any = None,
     value_fn: Optional[Callable[[Any], Any]] = None,
-    skip_notification: Optional[bool] = None) -> Any:
+    skip_notification: Optional[bool] = None,
+) -> Any:
   """Recursive patch values on condition.
 
   Args:

pyglove/core/patching/rule_based.py

@@ -16,9 +16,9 @@
 import re
 import typing
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
-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 Patcher(symbolic.Functor):
@@ -181,7 +181,9 @@ _PATCHER_REGISTRY = _PatcherRegistry()
 
 def patcher(
     args: Optional[List[Tuple[str, pg_typing.ValueSpec]]] = None,
-    name: Optional[str] = None) -> Any:
+    name: Optional[str] = None,
+    auto_typing: bool = False,
+    auto_doc: bool = False) -> Any:
   """Decorate a function into a Patcher and register it.
 
   A patcher function is defined as:
@@ -209,11 +211,20 @@ def patcher(
     args: A list of (arg_name, arg_value_spec) to schematize patcher arguments.
     name: String to be used as patcher name in URI. If None, function name will
       be used as patcher name.
+    auto_typing: If True, automatically inference the typing from the
+      function signature.
+    auto_doc: If True, use the docstring of the function as the docstring of
+      the patcher.
 
   Returns:
     A decorator that converts a function into a Patcher subclass.
   """
-  functor_decorator = symbolic.functor(args, base_class=Patcher)
+  functor_decorator = symbolic.functor(
+      args,
+      base_class=Patcher,
+      auto_typing=auto_typing,
+      auto_doc=auto_doc
+  )
   def _decorator(fn):
     """Returns decorated Patcher class."""
     cls = functor_decorator(fn)
@@ -339,7 +350,7 @@ def from_uri(uri: str) -> Patcher:
   name, args, kwargs = parse_uri(uri)
   patcher_cls = typing.cast(Type[Any], _PATCHER_REGISTRY.get(name))
   args, kwargs = parse_args(patcher_cls.__signature__, args, kwargs)
-  return patcher_cls(object_utils.MISSING_VALUE, *args, **kwargs)
+  return patcher_cls(utils.MISSING_VALUE, *args, **kwargs)
 
 
 def parse_uri(uri: str) -> Tuple[str, List[str], Dict[str, str]]:
@@ -456,7 +467,8 @@ def parse_arg(patcher_id: str, arg_name: str,
         f'{value_spec!r} cannot be used for Patcher argument.\n'
         f'Consider to treat this argument as string and parse it yourself.')
   return value_spec.apply(
-      arg, root_path=object_utils.KeyPath.parse(f'{patcher_id}.{arg_name}'))
+      arg, root_path=utils.KeyPath.parse(f'{patcher_id}.{arg_name}')
+  )
 
 
 def parse_list(string: str,

pyglove/core/patching/rule_based_test.py

@@ -48,6 +48,18 @@ class PatcherTest(unittest.TestCase):
         TypeError, 'The 1st argument of .* must be a symbolic type'):
       set_value1(k='a').patch(1)  # pylint: disable=not-callable, no-value-for-parameter
 
+  def test_patcher_with_auto_typing(self):
+    @rule_based.patcher(auto_typing=True)
+    def set_value2(unused_src, k: str, v: int):
+      return {
+          k: v
+      }
+    self.assert_patch_equal(symbolic.Dict(), set_value2(k='a', v=1), {'a': 1})  # pylint: disable=no-value-for-parameter
+    self.assert_patch_equal(symbolic.Dict(), 'set_value2?a&1', {'a': 1})
+    self.assert_patch_equal(symbolic.Dict(), 'set_value2?a&v=1', {'a': 1})
+    self.assert_patch_equal(symbolic.Dict(), 'set_value2?k=a&v=1', {'a': 1})
+    self.assertIn('set_value2', rule_based.patcher_names())
+
   def test_patcher_with_typing(self):
     @rule_based.patcher([
         ('k', pg_typing.Str()),
@@ -375,10 +387,21 @@ class PatcherHelpersTest(unittest.TestCase):
     signature = pg_typing.Signature(
         pg_typing.CallableType.FUNCTION, 'foo', '__main__',
         [
-            pg_typing.Argument('src', pg_typing.Any()),
-            pg_typing.Argument('x', pg_typing.Int()),
-            pg_typing.Argument('y', pg_typing.List(
-                pg_typing.Float(min_value=0.0, max_value=1.0)))
+            pg_typing.Argument(
+                'src',
+                pg_typing.Argument.Kind.POSITIONAL_OR_KEYWORD,
+                pg_typing.Any()
+            ),
+            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.List(pg_typing.Float(min_value=0.0, max_value=1.0))
+            )
         ])
     args, kwargs = rule_based.parse_args(signature, ['0'], {'y': '0.1:0.5'})
     self.assertEqual(args, [0])

pyglove/core/symbolic/__init__.py

@@ -95,6 +95,7 @@ from pyglove.core.symbolic.inferred import ValueFromParentChain
 # Reference type.
 from pyglove.core.symbolic.ref import Ref
 from pyglove.core.symbolic.ref import maybe_ref
+from pyglove.core.symbolic.ref import deref
 
 # Symbolic operations.
 from pyglove.core.symbolic.base import traverse
@@ -119,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
@@ -140,11 +142,4 @@ from pyglove.core.symbolic.list import mark_as_insertion
 # Error types.
 from pyglove.core.symbolic.base import WritePermissionError
 
-# TODO(daiyip): internal dependencies, remove later.
-from pyglove.core.symbolic.schema_utils import formalize_schema
-from pyglove.core.symbolic.schema_utils import augment_schema
-from pyglove.core.symbolic.schema_utils import function_schema
-from pyglove.core.symbolic.schema_utils import update_schema
-
-
 # pylint: enable=g-bad-import-order