bead 0.1.0__py3-none-any.whl

bead/data/repository.py

@@ -0,0 +1,358 @@
+"""Repository pattern for data access with optional caching.
+
+This module provides a generic Repository class that implements CRUD operations
+for Pydantic models, with optional in-memory caching for efficient access.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from uuid import UUID
+
+from pydantic import BaseModel
+
+from bead.data.serialization import (
+    append_jsonlines,
+    read_jsonlines,
+    write_jsonlines,
+)
+
+
+class Repository[T: BaseModel]:
+    """Generic repository for CRUD operations on Pydantic models.
+
+    Provides create, read, update, delete operations with JSONLines file storage
+    and optional in-memory caching for efficient data access.
+
+    Type Parameters
+    ---------------
+    T : BaseModel
+        Pydantic model type this repository manages
+
+    Parameters
+    ----------
+    model_class : type[T]
+        The Pydantic model class this repository manages
+    storage_path : Path
+        Path to the JSONLines file for persistent storage
+    use_cache : bool, optional
+        Whether to use in-memory caching (default: True)
+
+    Attributes
+    ----------
+    model_class : type[T]
+        The Pydantic model class
+    storage_path : Path
+        Path to storage file
+    use_cache : bool
+        Whether caching is enabled
+    cache : dict[UUID, T]
+        In-memory cache of objects by ID
+
+    Examples
+    --------
+    >>> from pathlib import Path
+    >>> from bead.data.base import BeadBaseModel
+    >>> class MyModel(BeadBaseModel):
+    ...     name: str
+    >>> repo = Repository[MyModel](
+    ...     model_class=MyModel,
+    ...     storage_path=Path("data/models.jsonl"),
+    ...     use_cache=True
+    ... )
+    >>> obj = MyModel(name="test")
+    >>> repo.add(obj)
+    >>> loaded = repo.get(obj.id)
+    >>> loaded.name
+    'test'
+    >>> repo.count()
+    1
+    """
+
+    def __init__(
+        self, model_class: type[T], storage_path: Path, use_cache: bool = True
+    ) -> None:
+        self.model_class = model_class
+        self.storage_path = storage_path
+        self.use_cache = use_cache
+        self.cache: dict[UUID, T] = {}
+
+        # load cache on init if enabled and file exists
+        if self.use_cache and self.storage_path.exists():
+            self._load_cache()
+
+    def _load_cache(self) -> None:
+        """Load all objects from storage into cache.
+
+        Called during initialization if caching is enabled and the storage
+        file exists.
+        """
+        objects = read_jsonlines(self.storage_path, self.model_class)
+        self.cache = {obj.id: obj for obj in objects}  # type: ignore[attr-defined]
+
+    def get(self, object_id: UUID) -> T | None:
+        """Get object by ID.
+
+        Parameters
+        ----------
+        object_id
+            ID of the object to retrieve.
+
+        Returns
+        -------
+        T | None
+            The object if found, None otherwise.
+
+        Examples
+        --------
+        >>> repo = Repository[MyModel](MyModel, Path("data.jsonl"))
+        >>> obj = MyModel(name="test")
+        >>> repo.add(obj)
+        >>> loaded = repo.get(obj.id)
+        >>> loaded is not None
+        True
+        """
+        if self.use_cache:
+            return self.cache.get(object_id)
+        else:
+            # scan file for object
+            if not self.storage_path.exists():
+                return None
+            objects = read_jsonlines(self.storage_path, self.model_class)
+            for obj in objects:
+                if obj.id == object_id:  # type: ignore[attr-defined]
+                    return obj
+            return None
+
+    def get_all(self) -> list[T]:
+        """Get all objects.
+
+        Returns
+        -------
+        list[T]
+            List of all objects in the repository
+
+        Examples
+        --------
+        >>> repo = Repository[MyModel](MyModel, Path("data.jsonl"))
+        >>> repo.add(MyModel(name="test1"))
+        >>> repo.add(MyModel(name="test2"))
+        >>> len(repo.get_all())
+        2
+        """
+        if self.use_cache:
+            return list(self.cache.values())
+        else:
+            if not self.storage_path.exists():
+                return []
+            return read_jsonlines(self.storage_path, self.model_class)
+
+    def add(self, obj: T) -> None:
+        """Add single object to repository.
+
+        Appends the object to the storage file and updates cache if enabled.
+
+        Parameters
+        ----------
+        obj
+            Object to add.
+
+        Examples
+        --------
+        >>> repo = Repository[MyModel](MyModel, Path("data.jsonl"))
+        >>> obj = MyModel(name="test")
+        >>> repo.add(obj)
+        >>> repo.exists(obj.id)
+        True
+        """
+        # create parent directories if needed
+        self.storage_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # append to file
+        append_jsonlines([obj], self.storage_path)
+
+        # update cache
+        if self.use_cache:
+            self.cache[obj.id] = obj  # type: ignore[attr-defined]
+
+    def add_many(self, objects: list[T]) -> None:
+        """Add multiple objects to repository.
+
+        Appends all objects to the storage file and updates cache if enabled.
+
+        Parameters
+        ----------
+        objects
+            List of objects to add.
+
+        Examples
+        --------
+        >>> repo = Repository[MyModel](MyModel, Path("data.jsonl"))
+        >>> objs = [MyModel(name="test1"), MyModel(name="test2")]
+        >>> repo.add_many(objs)
+        >>> repo.count()
+        2
+        """
+        if not objects:
+            return
+
+        # create parent directories if needed
+        self.storage_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # append to file
+        append_jsonlines(objects, self.storage_path)
+
+        # update cache
+        if self.use_cache:
+            for obj in objects:
+                self.cache[obj.id] = obj  # type: ignore[attr-defined]
+
+    def update(self, obj: T) -> None:
+        """Update existing object.
+
+        Rewrites the entire storage file with the updated object.
+
+        Parameters
+        ----------
+        obj
+            Object to update (must have existing ID).
+
+        Examples
+        --------
+        >>> repo = Repository[MyModel](MyModel, Path("data.jsonl"))
+        >>> obj = MyModel(name="test")
+        >>> repo.add(obj)
+        >>> obj.name = "updated"
+        >>> repo.update(obj)
+        >>> loaded = repo.get(obj.id)
+        >>> loaded.name
+        'updated'
+        """
+        # update in cache
+        if self.use_cache:
+            self.cache[obj.id] = obj  # type: ignore[attr-defined]
+
+        # rewrite file
+        objects = list(self.cache.values()) if self.use_cache else self.get_all()
+        # replace the object in the list
+        objects = [o if o.id != obj.id else obj for o in objects]  # type: ignore[attr-defined]
+
+        # create parent directories if needed
+        self.storage_path.parent.mkdir(parents=True, exist_ok=True)
+
+        write_jsonlines(objects, self.storage_path)
+
+    def delete(self, object_id: UUID) -> None:
+        """Delete object by ID.
+
+        Rewrites the entire storage file without the deleted object.
+
+        Parameters
+        ----------
+        object_id
+            ID of object to delete.
+
+        Examples
+        --------
+        >>> repo = Repository[MyModel](MyModel, Path("data.jsonl"))
+        >>> obj = MyModel(name="test")
+        >>> repo.add(obj)
+        >>> repo.delete(obj.id)
+        >>> repo.exists(obj.id)
+        False
+        """
+        # remove from cache
+        if self.use_cache:
+            self.cache.pop(object_id, None)
+
+        # rewrite file without the object
+        objects = list(self.cache.values()) if self.use_cache else self.get_all()
+        objects = [o for o in objects if o.id != object_id]  # type: ignore[attr-defined]
+
+        if objects:
+            self.storage_path.parent.mkdir(parents=True, exist_ok=True)
+            write_jsonlines(objects, self.storage_path)
+        elif self.storage_path.exists():
+            # if no objects left, delete the file
+            self.storage_path.unlink()
+
+    def exists(self, object_id: UUID) -> bool:
+        """Check if object exists.
+
+        Parameters
+        ----------
+        object_id
+            ID of object to check.
+
+        Returns
+        -------
+        bool
+            True if object exists, False otherwise.
+
+        Examples
+        --------
+        >>> repo = Repository[MyModel](MyModel, Path("data.jsonl"))
+        >>> obj = MyModel(name="test")
+        >>> repo.add(obj)
+        >>> repo.exists(obj.id)
+        True
+        """
+        return self.get(object_id) is not None
+
+    def count(self) -> int:
+        """Count objects in repository.
+
+        Returns
+        -------
+        int
+            Number of objects
+
+        Examples
+        --------
+        >>> repo = Repository[MyModel](MyModel, Path("data.jsonl"))
+        >>> repo.count()
+        0
+        >>> repo.add(MyModel(name="test"))
+        >>> repo.count()
+        1
+        """
+        if self.use_cache:
+            return len(self.cache)
+        else:
+            if not self.storage_path.exists():
+                return 0
+            return len(read_jsonlines(self.storage_path, self.model_class))
+
+    def clear(self) -> None:
+        """Clear all objects and delete storage file.
+
+        Examples
+        --------
+        >>> repo = Repository[MyModel](MyModel, Path("data.jsonl"))
+        >>> repo.add(MyModel(name="test"))
+        >>> repo.clear()
+        >>> repo.count()
+        0
+        """
+        # clear cache
+        self.cache.clear()
+
+        # delete file
+        if self.storage_path.exists():
+            self.storage_path.unlink()
+
+    def rebuild_cache(self) -> None:
+        """Rebuild cache from storage.
+
+        Reloads all objects from storage into the cache. Useful if the storage
+        file was modified externally.
+
+        Examples
+        --------
+        >>> repo = Repository[MyModel](MyModel, Path("data.jsonl"), use_cache=True)
+        >>> repo.rebuild_cache()
+        """
+        if not self.storage_path.exists():
+            self.cache.clear()
+        else:
+            self._load_cache()

bead/data/serialization.py

@@ -0,0 +1,249 @@
+"""JSONLines serialization utilities for bead package.
+
+This module provides functions for reading, writing, streaming, and appending
+Pydantic models to/from JSONLines format files. JSONLines is a convenient format
+for storing multiple JSON objects, with one object per line.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, ValidationError
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+class SerializationError(Exception):
+    """Exception raised when serialization to JSONLines fails.
+
+    This exception is raised when writing Pydantic objects to JSONLines
+    format encounters an error, such as file I/O issues or validation failures.
+    """
+
+    pass
+
+
+class DeserializationError(Exception):
+    """Exception raised when deserialization from JSONLines fails.
+
+    This exception is raised when reading JSONLines format into Pydantic objects
+    encounters an error, such as file not found, invalid JSON, or validation failures.
+    """
+
+    pass
+
+
+def write_jsonlines[T: BaseModel](
+    objects: Sequence[T],
+    path: Path | str,
+    validate: bool = True,
+    append: bool = False,
+) -> None:
+    """Write Pydantic objects to JSONLines file.
+
+    Serializes a sequence of Pydantic model instances to a JSONLines file,
+    with one JSON object per line. Each object is validated before writing
+    if validate=True.
+
+    Parameters
+    ----------
+    objects
+        Sequence of Pydantic model instances to serialize.
+    path
+        Path to the output file.
+    validate
+        Whether to validate objects before writing (default: True).
+    append
+        Whether to append to existing file or overwrite (default: False).
+
+    Raises
+    ------
+    SerializationError
+        If writing fails due to I/O error or validation failure
+
+    Examples
+    --------
+    >>> from pathlib import Path
+    >>> from bead.data.base import BeadBaseModel
+    >>> class TestModel(BeadBaseModel):
+    ...     name: str
+    >>> objects = [TestModel(name="test1"), TestModel(name="test2")]
+    >>> write_jsonlines(objects, Path("output.jsonl"))  # doctest: +SKIP
+    """
+    path = Path(path)
+    mode = "a" if append else "w"
+
+    try:
+        with path.open(mode, encoding="utf-8") as f:
+            for obj in objects:
+                # model_dump_json() handles validation
+                json_str = obj.model_dump_json()
+                f.write(json_str + "\n")
+    except (OSError, ValidationError) as e:
+        raise SerializationError(f"Failed to write to {path}: {e}") from e
+
+
+def read_jsonlines[T: BaseModel](
+    path: Path | str,
+    model_class: type[T],
+    validate: bool = True,
+    skip_errors: bool = False,
+) -> list[T]:
+    """Read JSONLines file into list of Pydantic objects.
+
+    Deserializes a JSONLines file into a list of Pydantic model instances.
+    Each line should contain a valid JSON object. Empty lines are skipped.
+
+    Parameters
+    ----------
+    path
+        Path to the input file.
+    model_class
+        Pydantic model class to deserialize into.
+    validate
+        Whether to validate objects during parsing (default: True).
+    skip_errors
+        Whether to skip invalid lines or raise error (default: False).
+
+    Returns
+    -------
+    list[T]
+        List of deserialized Pydantic objects
+
+    Raises
+    ------
+    DeserializationError
+        If reading fails due to file not found, invalid JSON, or validation failure
+        (unless skip_errors=True)
+
+    Examples
+    --------
+    >>> from pathlib import Path
+    >>> from bead.data.base import BeadBaseModel
+    >>> class TestModel(BeadBaseModel):
+    ...     name: str
+    >>> objects = read_jsonlines(Path("input.jsonl"), TestModel)  # doctest: +SKIP
+    """
+    path = Path(path)
+    objects: list[T] = []
+
+    try:
+        with path.open("r", encoding="utf-8") as f:
+            for line_num, line in enumerate(f, start=1):
+                line = line.strip()
+                if not line:  # skip empty lines
+                    continue
+
+                try:
+                    obj = model_class.model_validate_json(line)
+                    objects.append(obj)
+                except ValidationError as e:
+                    if skip_errors:
+                        continue
+                    raise DeserializationError(
+                        f"Failed to parse line {line_num} in {path}: {e}"
+                    ) from e
+    except OSError as e:
+        raise DeserializationError(f"Failed to read from {path}: {e}") from e
+
+    return objects
+
+
+def stream_jsonlines[T: BaseModel](
+    path: Path | str,
+    model_class: type[T],
+    validate: bool = True,
+) -> Iterator[T]:
+    """Stream JSONLines file as iterator of Pydantic objects.
+
+    Memory-efficient iterator that yields Pydantic model instances one at a time
+    from a JSONLines file. Useful for processing large files without loading
+    everything into memory.
+
+    Parameters
+    ----------
+    path
+        Path to the input file.
+    model_class
+        Pydantic model class to deserialize into.
+    validate
+        Whether to validate objects during parsing (default: True).
+
+    Yields
+    ------
+    T
+        Pydantic model instances one at a time.
+
+    Raises
+    ------
+    DeserializationError
+        If reading fails due to file not found, invalid JSON, or validation failure
+
+    Examples
+    --------
+    >>> from pathlib import Path
+    >>> from bead.data.base import BeadBaseModel
+    >>> class TestModel(BeadBaseModel):
+    ...     name: str
+    >>> for obj in stream_jsonlines(Path("input.jsonl"), TestModel):  # doctest: +SKIP
+    ...     print(obj.name)
+    """
+    path = Path(path)
+
+    try:
+        with path.open("r", encoding="utf-8") as f:
+            for line_num, line in enumerate(f, start=1):
+                line = line.strip()
+                if not line:  # skip empty lines
+                    continue
+
+                try:
+                    obj = model_class.model_validate_json(line)
+                    yield obj
+                except ValidationError as e:
+                    raise DeserializationError(
+                        f"Failed to parse line {line_num} in {path}: {e}"
+                    ) from e
+    except OSError as e:
+        raise DeserializationError(f"Failed to read from {path}: {e}") from e
+
+
+def append_jsonlines[T: BaseModel](
+    objects: Sequence[T],
+    path: Path | str,
+    validate: bool = True,
+) -> None:
+    """Append Pydantic objects to existing JSONLines file.
+
+    Convenience wrapper around write_jsonlines with append=True. Adds objects
+    to the end of an existing JSONLines file, or creates a new file if it
+    doesn't exist.
+
+    Parameters
+    ----------
+    objects
+        Sequence of Pydantic model instances to serialize.
+    path
+        Path to the output file.
+    validate
+        Whether to validate objects before writing (default: True).
+
+    Raises
+    ------
+    SerializationError
+        If appending fails due to I/O error or validation failure
+
+    Examples
+    --------
+    >>> from pathlib import Path
+    >>> from bead.data.base import BeadBaseModel
+    >>> class TestModel(BeadBaseModel):
+    ...     name: str
+    >>> objects = [TestModel(name="test3"), TestModel(name="test4")]
+    >>> append_jsonlines(objects, Path("output.jsonl"))  # doctest: +SKIP
+    """
+    write_jsonlines(objects, path, validate=validate, append=True)

bead/data/timestamps.py

@@ -0,0 +1,89 @@
+"""ISO 8601 timestamp utilities for bead package.
+
+This module provides functions for creating, parsing, and formatting ISO 8601
+timestamps with timezone information. All timestamps use UTC timezone.
+"""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+
+
+def now_iso8601() -> datetime:
+    """Get current UTC datetime with timezone information.
+
+    Returns the current time in UTC with timezone info attached. This is
+    preferred over datetime.utcnow() which is deprecated and doesn't include
+    timezone information.
+
+    Returns
+    -------
+    datetime
+        Current UTC datetime with timezone information
+
+    Examples
+    --------
+    >>> dt = now_iso8601()
+    >>> dt.tzinfo is not None
+    True
+    >>> dt.tzinfo == UTC
+    True
+    """
+    return datetime.now(UTC)
+
+
+def parse_iso8601(timestamp: str) -> datetime:
+    """Parse ISO 8601 timestamp string to datetime.
+
+    Parses an ISO 8601 formatted string into a datetime object. The string
+    should include timezone information.
+
+    Parameters
+    ----------
+    timestamp
+        ISO 8601 formatted timestamp string (e.g., "2025-10-17T14:23:45.123456+00:00").
+
+    Returns
+    -------
+    datetime
+        Parsed datetime with timezone information
+
+    Examples
+    --------
+    >>> dt_str = "2025-10-17T14:23:45.123456+00:00"
+    >>> dt = parse_iso8601(dt_str)
+    >>> dt.year
+    2025
+    >>> dt.month
+    10
+    """
+    return datetime.fromisoformat(timestamp)
+
+
+def format_iso8601(dt: datetime) -> str:
+    """Format datetime as ISO 8601 string.
+
+    Converts a datetime object to an ISO 8601 formatted string. If the datetime
+    doesn't have timezone information, it will be assumed to be UTC.
+
+    Parameters
+    ----------
+    dt
+        Datetime to format.
+
+    Returns
+    -------
+    str
+        ISO 8601 formatted string
+
+    Examples
+    --------
+    >>> dt = now_iso8601()
+    >>> formatted = format_iso8601(dt)
+    >>> "+00:00" in formatted or "Z" in formatted
+    True
+    """
+    # if datetime is naive (no timezone), assume UTC
+    if dt.tzinfo is None:
+        dt = dt.replace(tzinfo=UTC)
+    return dt.isoformat()