krons 0.1.0__py3-none-any.whl

kronos/core/pile.py

@@ -0,0 +1,575 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import contextlib
+import threading
+from collections.abc import Callable, Iterator
+from typing import Any, Generic, Literal, TypeVar, overload
+from uuid import UUID
+
+from pydantic import Field, PrivateAttr, field_serializer, field_validator
+from typing_extensions import override
+
+from kronos.errors import ExistsError, NotFoundError
+from kronos.protocols import Containable, Deserializable, Serializable, implements
+from kronos.types import Unset, UnsetType, is_unset
+from kronos.utils import extract_types, load_type_from_string, synchronized
+from kronos.utils.concurrency import Lock as AsyncLock
+
+from .element import Element
+from .progression import Progression
+
+__all__ = ("Pile",)
+
+T = TypeVar("T", bound=Element)
+
+
+@implements(
+    Containable,
+    Serializable,
+    Deserializable,
+)
+class Pile(Element, Generic[T]):
+    """Thread-safe typed collection with rich query interface.
+
+    A Pile is an ordered, type-validated container for Element subclasses.
+    Maintains insertion order via Progression, supports concurrent access.
+
+    Type-dispatched __getitem__:
+        pile[uuid|str]     -> T (single item by ID)
+        pile[int]          -> T (single item by index)
+        pile[slice]        -> Pile[T] (range)
+        pile[list|tuple]   -> Pile[T] (indices or UUIDs)
+        pile[Progression]  -> Pile[T] (ordered subset)
+        pile[callable]     -> Pile[T] (filter function)
+
+    Example:
+        pile = Pile[Node](items=[n1, n2], item_type=Node)
+        pile.add(n3)
+        filtered = pile[lambda x: x.metadata.get("active")]
+    """
+
+    _items: dict[UUID, T] = PrivateAttr(default_factory=dict)
+    _progression: Progression = PrivateAttr(default_factory=Progression)
+    _lock: threading.RLock = PrivateAttr(default_factory=threading.RLock)
+    _async_lock: AsyncLock = PrivateAttr(default_factory=AsyncLock)
+
+    @property
+    def progression(self) -> Progression:
+        """Read-only copy of progression order."""
+        return Progression(order=list(self._progression.order), name=self._progression.name)
+
+    item_type: set[type] | None = Field(
+        default=None,
+        frozen=True,
+        description="Allowed types for validation (None = any Element subclass)",
+    )
+    strict_type: bool = Field(
+        default=False,
+        frozen=True,
+        description="Enforce exact type match (disallow subclasses)",
+    )
+
+    @field_validator("item_type", mode="before")
+    @classmethod
+    def _normalize_item_type(cls, v: Any) -> set[type] | None:
+        """Normalize item_type to set[type] from various input formats."""
+        if v is None:
+            return None
+        if isinstance(v, list) and v and isinstance(v[0], str):
+            return {load_type_from_string(type_str) for type_str in v}
+        return extract_types(v)
+
+    @override
+    def __init__(
+        self,
+        items: list[T] | None = None,
+        item_type: type[T] | set[type] | list[type] | None = None,
+        order: list[UUID] | Progression | None = None,
+        strict_type: bool = False,
+        **kwargs,
+    ):
+        """Initialize Pile with optional items.
+
+        Args:
+            items: Initial items to add
+            item_type: Type(s) for validation (type, set, list, or Union)
+            order: Custom order (list of UUIDs or Progression)
+            strict_type: Enforce exact type match (no subclasses)
+            **kwargs: Element fields (id, created_at, metadata)
+
+        Raises:
+            NotFoundError: If order contains UUID not in items
+            TypeError: If item type validation fails
+        """
+        super().__init__(**{"item_type": item_type, "strict_type": strict_type, **kwargs})
+
+        if items:
+            for item in items:
+                self.add(item)
+
+        if order:
+            order_list = list(order.order) if isinstance(order, Progression) else order
+            for uid in order_list:
+                if uid not in self._items:
+                    raise NotFoundError(f"UUID {uid} in order not found in items")
+            self._progression = Progression(order=order_list)
+
+    @field_serializer("item_type")
+    def _serialize_item_type(self, v: set[type] | None) -> list[str] | None:
+        """Serialize item_type to list of fully-qualified type names."""
+        if v is None:
+            return None
+        return [f"{t.__module__}.{t.__name__}" for t in v]
+
+    @override
+    def to_dict(
+        self,
+        mode: Literal["python", "json", "db"] = "python",
+        created_at_format: (Literal["datetime", "isoformat", "timestamp"] | UnsetType) = Unset,
+        meta_key: str | UnsetType = Unset,
+        item_meta_key: str | UnsetType = Unset,
+        item_created_at_format: (Literal["datetime", "isoformat", "timestamp"] | UnsetType) = Unset,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """Serialize pile with items in progression order.
+
+        Args:
+            mode: Serialization mode (python/json/db)
+            created_at_format: Timestamp format for Pile itself
+            meta_key: Rename Pile metadata field
+            item_meta_key: Metadata key name for items
+            item_created_at_format: Timestamp format for items
+            **kwargs: Passed to model_dump()
+
+        Returns:
+            Dict with pile fields and serialized items in order
+        """
+        data = super().to_dict(
+            mode=mode, created_at_format=created_at_format, meta_key=meta_key, **kwargs
+        )
+
+        actual_meta_key = (
+            meta_key
+            if not is_unset(meta_key)
+            else ("node_metadata" if mode == "db" else "metadata")
+        )
+
+        if self._progression.name and actual_meta_key in data:
+            data[actual_meta_key]["progression_name"] = self._progression.name
+
+        item_mode = "python" if mode == "python" else "json"
+        data["items"] = [
+            i.to_dict(
+                mode=item_mode,
+                meta_key=item_meta_key,
+                created_at_format=item_created_at_format,
+            )
+            for i in self
+        ]
+
+        return data
+
+    # ==================== Core Operations ====================
+
+    @synchronized
+    def add(self, item: T) -> None:
+        """Add item to pile. Raises ExistsError if duplicate, TypeError if invalid type."""
+        self._validate_type(item)
+
+        if item.id in self._items:
+            raise ExistsError(f"Item {item.id} already exists in pile")
+
+        self._items[item.id] = item
+        self._progression.append(item.id)
+
+    @synchronized
+    def remove(self, item_id: UUID | str | Element) -> T:
+        """Remove and return item. Raises NotFoundError if not found."""
+        uid = self._coerce_id(item_id)
+
+        try:
+            item = self._items.pop(uid)
+        except KeyError:
+            raise NotFoundError(f"Item {uid} not found in pile") from None
+
+        self._progression.remove(uid)
+        return item
+
+    @synchronized
+    def pop(self, item_id: UUID | str | Element, default: Any = ...) -> T | Any:
+        """Remove and return item, or default if not found. Raises NotFoundError if no default."""
+        uid = self._coerce_id(item_id)
+
+        try:
+            item = self._items.pop(uid)
+            self._progression.remove(uid)
+            return item
+        except KeyError:
+            if default is ...:
+                raise NotFoundError(f"Item {uid} not found in pile") from None
+            return default
+
+    @synchronized
+    def get(self, item_id: UUID | str | Element, default: Any = ...) -> T | Any:
+        """Get item by ID, or default if not found. Raises NotFoundError if no default."""
+        uid = self._coerce_id(item_id)
+
+        try:
+            return self._items[uid]
+        except KeyError:
+            if default is ...:
+                raise NotFoundError(f"Item {uid} not found in pile") from None
+            return default
+
+    @synchronized
+    def update(self, item: T) -> None:
+        """Update existing item in-place. Raises NotFoundError if not found, TypeError if invalid."""
+        self._validate_type(item)
+
+        if item.id not in self._items:
+            raise NotFoundError(f"Item {item.id} not found in pile")
+
+        self._items[item.id] = item
+
+    @synchronized
+    def clear(self) -> None:
+        """Remove all items from pile."""
+        self._items.clear()
+        self._progression.clear()
+
+    # ==================== Set-like Operations ====================
+
+    @synchronized
+    def include(self, item: T) -> bool:
+        """Idempotent add: returns True if item is in pile after call, False on validation failure."""
+        if item.id in self._items:
+            return True
+        try:
+            self._validate_type(item)
+            self._items[item.id] = item
+            self._progression.append(item.id)
+            return True
+        except Exception:
+            return False
+
+    @synchronized
+    def exclude(self, item: UUID | str | Element) -> bool:
+        """Idempotent remove: returns True if item is not in pile after call, False on coercion failure."""
+        try:
+            uid = self._coerce_id(item)
+        except Exception:
+            return False
+        if uid not in self._items:
+            return True
+        self._items.pop(uid, None)
+        try:
+            self._progression.remove(uid)
+        except ValueError:
+            pass
+        return True
+
+    # ==================== Rich __getitem__ (Type Dispatch) ====================
+
+    @overload
+    def __getitem__(self, key: UUID | str) -> T: ...
+    @overload
+    def __getitem__(self, key: Progression) -> Pile[T]: ...
+    @overload
+    def __getitem__(self, key: int) -> T: ...
+    @overload
+    def __getitem__(self, key: slice) -> Pile[T]: ...
+    @overload
+    def __getitem__(self, key: list[int] | tuple[int, ...]) -> Pile[T]: ...
+    @overload
+    def __getitem__(self, key: list[UUID] | tuple[UUID, ...]) -> Pile[T]: ...
+    @overload
+    def __getitem__(self, key: Callable[[T], bool]) -> Pile[T]: ...
+
+    def __getitem__(self, key: Any) -> T | Pile[T]:
+        """Type-dispatched query: UUID/str/int -> T; slice/list/tuple/Progression/callable -> Pile[T]."""
+        if isinstance(key, (UUID, str)):
+            return self.get(key)
+        elif isinstance(key, int):
+            return self._get_by_index(key)
+        elif isinstance(key, Progression):
+            return self._filter_by_progression(key)
+        elif isinstance(key, slice):
+            return self._get_by_slice(key)
+        elif isinstance(key, (list, tuple)):
+            return self._get_by_list(key)
+        elif callable(key):
+            return self._filter_by_function(key)
+        else:
+            raise TypeError(
+                f"Invalid key type: {type(key)}. Expected UUID, str, int, slice, list, tuple, Progression, or callable"
+            )
+
+    def _filter_by_progression(self, prog: Progression) -> Pile[T]:
+        """Return new Pile with items in progression order. Raises NotFoundError if UUID missing."""
+        if any(uid not in self._items for uid in prog):
+            raise NotFoundError("Some items from progression not found in pile")
+
+        return Pile(
+            items=[self._items[uid] for uid in prog],
+            item_type=self.item_type,
+            strict_type=self.strict_type,
+        )
+
+    @synchronized
+    def _get_by_index(self, index: int) -> T:
+        """Get item by position in progression order."""
+        uid: UUID = self._progression[index]
+        return self._items[uid]
+
+    @synchronized
+    def _get_by_slice(self, s: slice) -> Pile[T]:
+        """Return new Pile with items from slice range."""
+        uids: list[UUID] = self._progression[s]
+        return Pile(
+            items=[self._items[uid] for uid in uids],
+            item_type=self.item_type,
+            strict_type=self.strict_type,
+        )
+
+    @synchronized
+    def _get_by_list(self, keys: list | tuple) -> Pile[T]:
+        """Return new Pile from list of indices or UUIDs. No mixing allowed."""
+        if not keys:
+            raise ValueError("Cannot get items with empty list/tuple")
+
+        first = keys[0]
+        if isinstance(first, int):
+            if not all(isinstance(k, int) for k in keys):
+                raise TypeError("Cannot mix int and UUID in list/tuple indexing")
+            items = [self._get_by_index(idx) for idx in keys]
+        elif isinstance(first, (UUID, str)):
+            if not all(isinstance(k, (UUID, str)) for k in keys):
+                raise TypeError("Cannot mix int and UUID in list/tuple indexing")
+            items = [self.get(uid) for uid in keys]
+        else:
+            raise TypeError(f"list/tuple must contain only int or UUID, got {type(first)}")
+
+        return Pile(
+            items=items,
+            item_type=self.item_type,
+            strict_type=self.strict_type,
+        )
+
+    def _filter_by_function(self, func: Callable[[T], bool]) -> Pile[T]:
+        """Return new Pile with items matching filter function."""
+        return Pile(
+            items=[item for item in self if func(item)],
+            item_type=self.item_type,
+            strict_type=self.strict_type,
+        )
+
+    def filter_by_type(self, item_type: type[T] | set[type] | list[type]) -> Pile[T]:
+        """Return new Pile with items matching specified type(s).
+
+        Args:
+            item_type: Type(s) to filter by
+
+        Returns:
+            New Pile containing only items of requested type(s)
+
+        Raises:
+            TypeError: If requested type incompatible with pile's item_type
+            NotFoundError: If no items match
+        """
+        types_to_filter = extract_types(item_type)
+
+        if self.item_type is not None:
+            if self.strict_type:
+                invalid_types = types_to_filter - self.item_type
+                if invalid_types:
+                    raise TypeError(
+                        f"Types {invalid_types} not allowed in pile (allowed: {self.item_type})"
+                    )
+            else:
+                for t in types_to_filter:
+                    is_compatible = any(
+                        issubclass(t, allowed) or issubclass(allowed, t)
+                        for allowed in self.item_type
+                    )
+                    if not is_compatible:
+                        raise TypeError(
+                            f"Type {t} not compatible with allowed types {self.item_type}"
+                        )
+
+        filtered_items = [
+            item for item in self if any(isinstance(item, t) for t in types_to_filter)
+        ]
+
+        if not filtered_items:
+            raise NotFoundError(f"No items of type(s) {types_to_filter} found in pile")
+
+        return Pile(
+            items=filtered_items,
+            item_type=self.item_type,
+            strict_type=self.strict_type,
+        )
+
+    # ==================== Context Managers ====================
+
+    async def __aenter__(self) -> Pile[T]:
+        """Acquire async lock for context manager."""
+        await self._async_lock.acquire()
+        return self
+
+    async def __aexit__(self, _exc_type, _exc_val, _exc_tb) -> None:
+        """Release async lock."""
+        self._async_lock.release()
+
+    # ==================== Query Operations ====================
+
+    @synchronized
+    def __contains__(self, item: UUID | str | Element) -> bool:
+        """Check if item exists in pile by ID."""
+        with contextlib.suppress(Exception):
+            uid = self._coerce_id(item)
+            return uid in self._items
+        return False
+
+    @synchronized
+    def __len__(self) -> int:
+        return len(self._items)
+
+    def __bool__(self) -> bool:
+        return len(self._items) > 0
+
+    @synchronized
+    def __iter__(self) -> Iterator[T]:  # type: ignore[override]
+        """Iterate items in progression order. LSP override: yields T, not field tuples."""
+        for uid in self._progression:
+            yield self._items[uid]
+
+    def keys(self) -> Iterator[UUID]:
+        """Iterate UUIDs in progression order."""
+        return iter(self._progression)
+
+    def items(self) -> Iterator[tuple[UUID, T]]:
+        """Iterate (UUID, item) pairs in progression order."""
+        for i in self:
+            yield (i.id, i)
+
+    def __list__(self) -> list[T]:
+        """Return items as list in progression order."""
+        return [i for i in self]
+
+    def is_empty(self) -> bool:
+        """Check if pile contains no items."""
+        return len(self._items) == 0
+
+    # ==================== Validation ====================
+
+    def _validate_type(self, item: T) -> None:
+        """Validate item against pile's type constraints. Raises TypeError on failure."""
+        if not isinstance(item, Element):
+            raise TypeError(f"Item must be Element subclass, got {type(item)}")
+
+        if self.item_type is not None:
+            item_type_actual = type(item)
+
+            if self.strict_type:
+                if item_type_actual not in self.item_type:
+                    raise TypeError(
+                        f"Item type {item_type_actual} not in allowed types {self.item_type} "
+                        "(strict_type=True, no subclasses allowed)"
+                    )
+            else:
+                if not any(issubclass(item_type_actual, t) for t in self.item_type):
+                    raise TypeError(
+                        f"Item type {item_type_actual} is not a subclass of any allowed type {self.item_type}"
+                    )
+
+    # ==================== Deserialization ====================
+
+    @classmethod
+    @override
+    def from_dict(
+        cls,
+        data: dict[str, Any],
+        meta_key: str | UnsetType = Unset,
+        item_meta_key: str | UnsetType = Unset,
+        **kwargs: Any,
+    ) -> Pile[T]:
+        """Deserialize Pile from dict. Validates all item types before deserializing.
+
+        Args:
+            data: Serialized pile data
+            meta_key: Restore metadata from this key (db mode compatibility)
+            item_meta_key: Metadata key for item deserialization
+            **kwargs: Additional arguments (e.g., item_type override)
+
+        Returns:
+            Reconstructed Pile with all items
+
+        Raises:
+            TypeError: If any item type incompatible with pile's constraints
+        """
+        from .element import Element
+
+        data = data.copy()
+
+        if not is_unset(meta_key) and meta_key in data:
+            data["metadata"] = data.pop(meta_key)
+        elif "node_metadata" in data and "metadata" not in data:
+            data["metadata"] = data.pop("node_metadata")
+        data.pop("node_metadata", None)
+
+        item_type_data = data.get("item_type") or kwargs.get("item_type")
+        strict_type = data.get("strict_type", False)
+
+        items_data = data.get("items", [])
+        if item_type_data is not None and items_data:
+            if (
+                isinstance(item_type_data, list)
+                and item_type_data
+                and isinstance(item_type_data[0], str)
+            ):
+                allowed_types = {load_type_from_string(type_str) for type_str in item_type_data}
+            else:
+                allowed_types = extract_types(item_type_data)
+
+            for item_dict in items_data:
+                kron_class = item_dict.get("metadata", {}).get("kron_class")
+                if kron_class:
+                    try:
+                        item_type_actual = load_type_from_string(kron_class)
+                    except ValueError:
+                        continue
+
+                    if strict_type:
+                        if item_type_actual not in allowed_types:
+                            raise TypeError(
+                                f"Item type {kron_class} not in allowed types {allowed_types} "
+                                "(strict_type=True)"
+                            )
+                    else:
+                        if not any(issubclass(item_type_actual, t) for t in allowed_types):
+                            raise TypeError(
+                                f"Item type {kron_class} is not a subclass of any allowed type {allowed_types}"
+                            )
+
+        pile_data = data.copy()
+        pile_data.pop("items", None)
+        pile_data.pop("item_type", None)
+        pile_data.pop("strict_type", None)
+        pile = cls(item_type=item_type_data, strict_type=strict_type, **pile_data)
+
+        metadata = data.get("metadata", {})
+        progression_name = metadata.get("progression_name")
+        if progression_name:
+            pile._progression.name = progression_name
+
+        for item_dict in items_data:
+            item = Element.from_dict(item_dict, meta_key=item_meta_key)
+            pile.add(item)  # type: ignore[arg-type]
+
+        return pile
+
+    def __repr__(self) -> str:
+        return f"Pile(len={len(self)})"