brainlessdb 0.1.0__py3-none-any.whl

brainless/entity.py

@@ -0,0 +1,572 @@
+"""Entity wrapper with attribute access and dirty tracking."""
+
+from __future__ import annotations
+
+import sys
+from dataclasses import fields, is_dataclass
+from typing import TYPE_CHECKING, Any, Callable, TypeVar, get_origin, get_type_hints
+
+if TYPE_CHECKING:
+    from brainless.collection import Collection
+
+T = TypeVar("T")
+L = TypeVar("L")
+K = TypeVar("K")
+V = TypeVar("V")
+
+# Cache for tracked dataclass subclasses
+_tracked_classes: dict[type, type] = {}
+# Cache for nested tracked dataclass subclasses
+_nested_tracked_classes: dict[type, type] = {}
+
+
+class TrackedList(list[L]):
+    """List that tracks mutations and syncs to parent.
+
+    When items are added, removed, or modified, the backing list in the
+    entity data is updated and mark_dirty is called.
+    """
+
+    def __init__(
+        self,
+        items: list[L],
+        backing_list: list[Any],
+        mark_dirty: Callable[[], None],
+        item_type: type | None = None,
+    ) -> None:
+        super().__init__(items)
+        self._backing_list = backing_list
+        self._mark_dirty = mark_dirty
+        self._item_type = item_type
+
+    def _to_dict(self, item: Any) -> Any:
+        """Convert item to dict for storage."""
+        if is_dataclass(item) and not isinstance(item, type):
+            from dataclasses import asdict
+            return asdict(item)
+        return item
+
+    def _sync_and_dirty(self) -> None:
+        """Sync to backing list and mark dirty."""
+        self._backing_list.clear()
+        self._backing_list.extend(self._to_dict(item) for item in list.__iter__(self))
+        self._mark_dirty()
+
+    def append(self, item: L) -> None:
+        super().append(item)
+        self._backing_list.append(self._to_dict(item))
+        self._mark_dirty()
+
+    def extend(self, items: list[L]) -> None:
+        super().extend(items)
+        self._backing_list.extend(self._to_dict(item) for item in items)
+        self._mark_dirty()
+
+    def insert(self, index: int, item: L) -> None:
+        super().insert(index, item)
+        self._backing_list.insert(index, self._to_dict(item))
+        self._mark_dirty()
+
+    def remove(self, item: L) -> None:
+        idx = self.index(item)
+        super().remove(item)
+        del self._backing_list[idx]
+        self._mark_dirty()
+
+    def pop(self, index: int = -1) -> L:
+        result = super().pop(index)
+        self._backing_list.pop(index)
+        self._mark_dirty()
+        return result
+
+    def clear(self) -> None:
+        super().clear()
+        self._backing_list.clear()
+        self._mark_dirty()
+
+    def __setitem__(self, index: int, item: L) -> None:  # type: ignore[override]
+        super().__setitem__(index, item)
+        self._backing_list[index] = self._to_dict(item)
+        self._mark_dirty()
+
+    def __delitem__(self, index: int) -> None:  # type: ignore[override]
+        super().__delitem__(index)
+        del self._backing_list[index]
+        self._mark_dirty()
+
+    def __iadd__(self, items: list[L]) -> "TrackedList[L]":  # type: ignore[override]
+        self.extend(items)
+        return self
+
+
+class TrackedDict(dict[K, V]):
+    """Dict that tracks mutations and syncs to parent.
+
+    When items are added, removed, or modified, the backing dict in the
+    entity data is updated and mark_dirty is called.
+    """
+
+    def __init__(
+        self,
+        items: dict[K, V],
+        backing_dict: dict[Any, Any],
+        mark_dirty: Callable[[], None],
+        value_type: type | None = None,
+    ) -> None:
+        super().__init__(items)
+        self._backing_dict = backing_dict
+        self._mark_dirty = mark_dirty
+        self._value_type = value_type
+
+    def _to_dict(self, item: Any) -> Any:
+        """Convert item to dict for storage."""
+        if is_dataclass(item) and not isinstance(item, type):
+            from dataclasses import asdict
+            return asdict(item)
+        return item
+
+    def _to_key(self, key: Any) -> Any:
+        """Convert key to storage format (stringify int keys for JSON)."""
+        return str(key) if isinstance(key, int) else key
+
+    def __setitem__(self, key: K, value: V) -> None:
+        super().__setitem__(key, value)
+        self._backing_dict[self._to_key(key)] = self._to_dict(value)
+        self._mark_dirty()
+
+    def __delitem__(self, key: K) -> None:
+        super().__delitem__(key)
+        del self._backing_dict[self._to_key(key)]
+        self._mark_dirty()
+
+    def pop(self, key: K, *args: Any) -> V:
+        result = super().pop(key, *args)
+        self._backing_dict.pop(self._to_key(key), None)
+        self._mark_dirty()
+        return result
+
+    def popitem(self) -> tuple[K, V]:
+        key, value = super().popitem()
+        self._backing_dict.pop(self._to_key(key), None)
+        self._mark_dirty()
+        return key, value
+
+    def clear(self) -> None:
+        super().clear()
+        self._backing_dict.clear()
+        self._mark_dirty()
+
+    def update(self, other: dict[K, V] | None = None, **kwargs: V) -> None:  # type: ignore[override]
+        if other:
+            super().update(other)
+            for k, v in other.items():
+                self._backing_dict[self._to_key(k)] = self._to_dict(v)
+        if kwargs:
+            super().update(kwargs)  # type: ignore[arg-type]
+            for k, v in kwargs.items():
+                self._backing_dict[self._to_key(k)] = self._to_dict(v)
+        self._mark_dirty()
+
+    def setdefault(self, key: K, default: V = None) -> V:  # type: ignore[assignment]
+        if key not in self:
+            self[key] = default
+        return self[key]
+
+
+def _get_hints(cls: type) -> dict[str, Any]:
+    """Get type hints for a class, resolving forward references."""
+    try:
+        # Get the module where the class is defined for proper resolution
+        module = sys.modules.get(cls.__module__, None)
+        globalns = getattr(module, "__dict__", {}) if module else {}
+        return get_type_hints(cls, globalns=globalns)
+    except NameError:
+        # Fall back to raw annotations and try to resolve manually
+        return _resolve_annotations(cls)
+
+
+def _resolve_annotations(cls: type) -> dict[str, Any]:
+    """Manually resolve string annotations by searching loaded modules."""
+    raw = getattr(cls, "__annotations__", {})
+    resolved: dict[str, Any] = {}
+
+    for name, hint in raw.items():
+        if isinstance(hint, str):
+            resolved[name] = _find_class(hint)
+        else:
+            resolved[name] = hint
+
+    return resolved
+
+
+def _find_class(name: str) -> Any:
+    """Find a class by name in loaded modules."""
+    # Handle generic types like list[Foo], Optional[Foo]
+    origin_match = name.split("[")[0].strip()
+    if origin_match in ("list", "List", "Optional", "Union"):
+        return name  # Can't resolve generics easily, return as string
+
+    # Search all loaded modules for the class
+    for module in sys.modules.values():
+        if module is None:
+            continue
+        cls = getattr(module, name, None)
+        if cls is not None and isinstance(cls, type):
+            return cls
+
+    # Not found, return as string
+    return name
+
+
+def _get_tracked_class(cls: type[T]) -> type[T]:
+    """Get or create a tracked subclass of a dataclass.
+
+    The tracked subclass syncs attribute changes back to the Entity.
+    """
+    if cls in _tracked_classes:
+        return _tracked_classes[cls]
+
+    class Tracked(cls):  # type: ignore[valid-type,misc]
+        __slots__ = ("_brainless_entity",)
+
+        def __setattr__(self, name: str, value: Any) -> None:
+            if name == "_brainless_entity":
+                object.__setattr__(self, name, value)
+                return
+
+            # Set on dataclass first
+            super().__setattr__(name, value)
+
+            # Sync to entity if tracked and field exists
+            entity = getattr(self, "_brainless_entity", None)
+            if entity is not None and name in entity._data:
+                entity[name] = value
+
+    Tracked.__name__ = f"Tracked{cls.__name__}"
+    Tracked.__qualname__ = f"Tracked{cls.__qualname__}"
+    _tracked_classes[cls] = Tracked
+    return Tracked  # type: ignore[return-value]
+
+
+def _get_nested_tracked_class(cls: type[T]) -> type[T]:
+    """Get or create a nested tracked subclass of a dataclass.
+
+    Nested tracked objects sync changes to their parent dict and call
+    a mark_dirty callback.
+    """
+    if cls in _nested_tracked_classes:
+        return _nested_tracked_classes[cls]
+
+    class NestedTracked(cls):  # type: ignore[valid-type,misc]
+        __slots__ = ("_brainless_parent_dict", "_brainless_mark_dirty")
+
+        def __setattr__(self, name: str, value: Any) -> None:
+            if name in ("_brainless_parent_dict", "_brainless_mark_dirty"):
+                object.__setattr__(self, name, value)
+                return
+
+            # Set on dataclass first
+            super().__setattr__(name, value)
+
+            # Sync to parent dict and mark dirty
+            parent_dict = getattr(self, "_brainless_parent_dict", None)
+            mark_dirty = getattr(self, "_brainless_mark_dirty", None)
+            if parent_dict is not None and name in parent_dict:
+                parent_dict[name] = value
+                if mark_dirty is not None:
+                    mark_dirty()
+
+    NestedTracked.__name__ = cls.__name__
+    NestedTracked.__qualname__ = cls.__qualname__
+    _nested_tracked_classes[cls] = NestedTracked
+    return NestedTracked  # type: ignore[return-value]
+
+
+def _make_nested_tracked(
+    cls: type[T],
+    data: dict[str, Any],
+    mark_dirty: Callable[[], None],
+) -> T:
+    """Create a nested tracked dataclass instance.
+
+    Changes sync to the provided data dict and call mark_dirty.
+    """
+    tracked_cls = _get_nested_tracked_class(cls)
+    hints = _get_hints(cls)
+
+    field_names = {f.name for f in fields(cls)}
+    kwargs: dict[str, Any] = {}
+
+    for name in field_names:
+        if name not in data:
+            continue
+        value = data[name]
+        field_type = hints.get(name)
+
+        # Recursively handle nested dataclasses
+        if field_type and is_dataclass(field_type) and isinstance(value, dict):
+            kwargs[name] = _make_nested_tracked(field_type, value, mark_dirty)
+        elif field_type and get_origin(field_type) is list:
+            args = getattr(field_type, "__args__", ())
+            if args and is_dataclass(args[0]) and isinstance(value, list):
+                tracked_items = [
+                    _make_nested_tracked(args[0], item, mark_dirty)
+                    if isinstance(item, dict) else item
+                    for item in value
+                ]
+                kwargs[name] = TrackedList(tracked_items, value, mark_dirty, args[0])
+            elif isinstance(value, list):
+                kwargs[name] = TrackedList(value.copy(), value, mark_dirty)
+            else:
+                kwargs[name] = value
+        elif field_type and get_origin(field_type) is dict:
+            args = getattr(field_type, "__args__", ())
+            if len(args) >= 2 and is_dataclass(args[1]) and isinstance(value, dict):
+                # dict[K, DataclassType] - track values
+                key_type = args[0]
+                val_type = args[1]
+                tracked_items = {}
+                for k, v in value.items():
+                    # Convert key back to original type if needed
+                    typed_key = int(k) if key_type is int and isinstance(k, str) else k
+                    tracked_items[typed_key] = (
+                        _make_nested_tracked(val_type, v, mark_dirty)
+                        if isinstance(v, dict) else v
+                    )
+                kwargs[name] = TrackedDict(tracked_items, value, mark_dirty, val_type)
+            elif isinstance(value, dict):
+                kwargs[name] = TrackedDict(value.copy(), value, mark_dirty)
+            else:
+                kwargs[name] = value
+        else:
+            kwargs[name] = value
+
+    instance = tracked_cls(**kwargs)
+    object.__setattr__(instance, "_brainless_parent_dict", data)
+    object.__setattr__(instance, "_brainless_mark_dirty", mark_dirty)
+    return instance
+
+
+def _make_tracked_instance(cls: type[T], entity: "Entity") -> T:
+    """Create a tracked dataclass instance linked to an Entity.
+
+    Changes to the instance's attributes are synced back to the Entity,
+    triggering dirty tracking and eventual persistence.
+    Handles nested dataclasses recursively.
+    """
+    tracked_cls = _get_tracked_class(cls)
+    hints = _get_hints(cls)
+
+    field_names = {f.name for f in fields(cls)}
+    kwargs: dict[str, Any] = {}
+
+    def mark_dirty() -> None:
+        entity.mark_dirty()
+
+    for name in field_names:
+        if name == "uuid":
+            kwargs["uuid"] = entity._uuid
+            continue
+        if name not in entity._data:
+            continue
+
+        value = entity._data[name]
+        field_type = hints.get(name)
+
+        # Recursively handle nested dataclasses
+        if field_type and is_dataclass(field_type) and isinstance(value, dict):
+            kwargs[name] = _make_nested_tracked(field_type, value, mark_dirty)
+        elif field_type and get_origin(field_type) is list:
+            args = getattr(field_type, "__args__", ())
+            if args and is_dataclass(args[0]) and isinstance(value, list):
+                tracked_items = [
+                    _make_nested_tracked(args[0], item, mark_dirty)
+                    if isinstance(item, dict) else item
+                    for item in value
+                ]
+                kwargs[name] = TrackedList(tracked_items, value, mark_dirty, args[0])
+            elif isinstance(value, list):
+                kwargs[name] = TrackedList(value.copy(), value, mark_dirty)
+            else:
+                kwargs[name] = value
+        elif field_type and get_origin(field_type) is dict:
+            args = getattr(field_type, "__args__", ())
+            if len(args) >= 2 and is_dataclass(args[1]) and isinstance(value, dict):
+                # dict[K, DataclassType] - track values
+                key_type = args[0]
+                val_type = args[1]
+                tracked_items = {}
+                for k, v in value.items():
+                    # Convert key back to original type if needed
+                    typed_key = int(k) if key_type is int and isinstance(k, str) else k
+                    tracked_items[typed_key] = (
+                        _make_nested_tracked(val_type, v, mark_dirty)
+                        if isinstance(v, dict) else v
+                    )
+                kwargs[name] = TrackedDict(tracked_items, value, mark_dirty, val_type)
+            elif isinstance(value, dict):
+                kwargs[name] = TrackedDict(value.copy(), value, mark_dirty)
+            else:
+                kwargs[name] = value
+        else:
+            kwargs[name] = value
+
+    instance = tracked_cls(**kwargs)
+    object.__setattr__(instance, "_brainless_entity", entity)
+    return instance
+
+
+def _convert_to_dataclass(data: dict[str, Any], cls: type[T], uuid: str | None = None) -> T:
+    """Recursively convert dict to dataclass, handling nested types.
+
+    If the dataclass has a 'uuid' field and uuid is provided, it will be
+    automatically populated.
+    """
+    if not is_dataclass(cls):
+        raise TypeError(f"Expected dataclass, got {cls.__name__}")
+
+    hints = _get_hints(cls)
+
+    converted: dict[str, Any] = {}
+
+    for field_name, field_type in hints.items():
+        if field_name not in data:
+            continue
+
+        value = data[field_name]
+        converted[field_name] = _convert_value(value, field_type)
+
+    # Auto-populate uuid if field exists in dataclass
+    if uuid is not None and "uuid" in hints and "uuid" not in converted:
+        converted["uuid"] = uuid
+
+    return cls(**converted)
+
+
+def _convert_value(value: Any, target_type: type) -> Any:
+    """Convert a single value to target type."""
+    if value is None:
+        return None
+
+    # String annotations (unresolved forward refs) - pass through
+    if isinstance(target_type, str):
+        return value
+
+    origin = get_origin(target_type)
+
+    # Handle list[SomeType]
+    if origin is list:
+        args = getattr(target_type, "__args__", ())
+        if args and is_dataclass(args[0]):
+            return [_convert_to_dataclass(item, args[0]) for item in value]
+        return value
+
+    # Handle nested dataclass
+    if is_dataclass(target_type) and isinstance(value, dict):
+        return _convert_to_dataclass(value, target_type)
+
+    return value
+
+
+class Entity:
+    """Wrapper around stored data with attribute access and dirty tracking
+
+    Provides dict-like and attribute access to fields. Automatically marks
+    the entity as dirty when values change, triggering background flush.
+    """
+
+    __slots__ = ("_collection", "_uuid", "_data", "_dirty")
+
+    def __init__(self, collection: Collection, uuid: str, data: dict[str, Any]) -> None:
+        object.__setattr__(self, "_collection", collection)
+        object.__setattr__(self, "_uuid", uuid)
+        object.__setattr__(self, "_data", data)
+        object.__setattr__(self, "_dirty", False)
+
+    @property
+    def uuid(self) -> str:
+        return self._uuid
+
+    @property
+    def data(self) -> dict[str, Any]:
+        """Raw entity data (excludes uuid)."""
+        return self._data
+
+    @property
+    def dirty(self) -> bool:
+        return self._dirty
+
+    def mark_dirty(self) -> None:
+        """Mark entity as modified, schedule for flush."""
+        if not self._dirty:
+            object.__setattr__(self, "_dirty", True)
+            self._collection.mark_dirty(self)
+
+    def mark_clean(self) -> None:
+        """Mark entity as persisted."""
+        object.__setattr__(self, "_dirty", False)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Return copy of entity data including uuid."""
+        return {"uuid": self._uuid, **self._data}
+
+    def as_type(self, cls: type[T]) -> T:
+        """Convert entity data to a dataclass instance
+
+        Handles nested dataclasses and lists of dataclasses.
+        If the dataclass has a 'uuid' field, it will be auto-populated.
+
+        @param cls: Target dataclass type
+        @return: Instance of cls populated with entity data
+        """
+        return _convert_to_dataclass(self._data, cls, uuid=self._uuid)
+
+    def __getattr__(self, name: str) -> Any:
+        if name.startswith("_"):
+            raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'")
+        try:
+            return self._data[name]
+        except KeyError:
+            raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'") from None
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        if name.startswith("_"):
+            object.__setattr__(self, name, value)
+            return
+
+        if name not in self._data:
+            raise AttributeError(f"Cannot add new field '{name}' to entity")
+
+        old_value = self._data[name]
+        if old_value != value:
+            self._data[name] = value
+            self._collection.on_field_change(self, name, old_value, value)
+            self.mark_dirty()
+
+    def __getitem__(self, key: str) -> Any:
+        if key == "uuid":
+            return self._uuid
+        return self._data[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        if key == "uuid":
+            raise KeyError("Cannot modify uuid")
+        if key not in self._data:
+            raise KeyError(f"Unknown field '{key}'")
+        old_value = self._data[key]
+        if old_value != value:
+            self._data[key] = value
+            self._collection.on_field_change(self, key, old_value, value)
+            self.mark_dirty()
+
+    def __contains__(self, key: str) -> bool:
+        return key == "uuid" or key in self._data
+
+    def __repr__(self) -> str:
+        dirty_marker = "*" if self._dirty else ""
+        # Show first few fields for context
+        fields = [f"{k}={v!r}" for k, v in list(self._data.items())[:3]]
+        fields_str = ", ".join(fields)
+        if len(self._data) > 3:
+            fields_str += ", ..."
+        return f"<{self._collection.name}:{self._uuid[-8:]}{dirty_marker} {fields_str}>"