grung-db/src/grung/types.py

from __future__ import annotations

from collections import namedtuple
from dataclasses import dataclass, field
from typing import Dict, List

import nanoid
from tinydb import TinyDB, where

from grung.exceptions import PointerReferenceError

Metadata = namedtuple("Metadata", ["table", "fields"])


@dataclass
class Field:
    """
    Represents a single field in a Record.
    """

    name: str
    value_type: type = str
    default: str = None
    unique: bool = False

    def serialize(self, rec: value_type, db: TinyDB) -> str:
        return str(rec)

    def deserialize(self, rec: value_type, db: TinyDB, recurse: bool = True) -> value_type:
        return rec


@dataclass
class Integer(Field):
    value_type: type = int
    default: int = 0

    def deserialize(self, rec: str, db: TinyDB, recurse: bool = True) -> value_type:
        return int(rec)


class Record(Dict[(str, Field)]):
    """
    Base type for a single database record.
    """

    def __init__(self, raw_doc: dict = {}, doc_id: int = None, **params):
        self.doc_id = doc_id
        fields = self.__class__.fields()
        self._metadata = Metadata(table=self.__class__.__name__, fields={f.name: f for f in fields})
        super().__init__(dict({field.name: field.default for field in fields}, **raw_doc, **params))

    @classmethod
    def fields(self):
        return [
            # 1% collision rate at ~2M records
            Field("uid", default=nanoid.generate(size=8), unique=True)
        ]

    def serialize(self, db):
        """
        Serialie every field on the record
        """
        rec = {}
        for name, _field in self._metadata.fields.items():
            rec[name] = _field.serialize(self[name], db)
        return self.__class__(rec, doc_id=self.doc_id)

    def deserialize(self, db, recurse: bool = True):
        """
        Deserialize every field on the record
        """
        rec = {}
        for name, _field in self._metadata.fields.items():
            rec[name] = _field.deserialize(self[name], db, recurse=recurse)
        return self.__class__(rec, doc_id=self.doc_id)

    def before_insert(self):
        pass

    def __setattr__(self, key, value):
        if key in self:
            self[key] = value
        super().__setattr__(key, value)

    def __getattr__(self, attr_name):
        if attr_name in self:
            return self.get(attr_name)
        return super().__getattr__(attr_name)

    def __hash__(self):
        return hash(str(dict(self)))

    def __repr__(self):
        return f"{self.__class__.__name__}[{self.doc_id}]: {self.items()}"


@dataclass
class Pointer(Field):
    """
    Store a string reference to a record.
    """

    name: str = ""
    value_type: type = Record

    def serialize(self, value: value_type, db: TinyDB) -> str:
        if value:
            if not value.doc_id:
                raise PointerReferenceError(value)
            return f"{value._metadata.table}::{value.uid}"
        return ""

    def deserialize(self, value: str, db: TinyDB, recurse=True) -> value_type:
        if value:
            pt, puid = value.split("::")
            if puid:
                return db.table(pt).search(where("uid") == puid, recurse=recurse)[0]
        return value


@dataclass
class Collection(Field):
    """
    A collection of pointers.
    """

    value_type: type = Record
    default: List[value_type] = field(default_factory=lambda: [])

    def _pointer(self, rec):
        return Pointer(value_type=type(rec))

    def serialize(self, values: List[Record], db: TinyDB) -> List[str]:
        if values:
            return [self._pointer(val).serialize(val, db=db) for val in values]

    def deserialize(self, values: List[str], db: TinyDB, recurse=True) -> List[value_type]:
        """
        Recursively deserialize the objects in this collection
        """
        if not recurse:
            return values
        if values:
            return [self._pointer(val).deserialize(val, db=db, recurse=recurse) for val in values]
        return []