grung-db/src/grung/types.py

from __future__ import annotations

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

import nanoid
from tinydb import where

from grung.exceptions import PointerReferenceError

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


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

    value_type = str

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

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

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


class Integer(Field):
    value_type = int

    default: value_type = 0

    def deserialize(self, rec: str, db: TinyDB) -> 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):
        # populate the metadata
        self._fields.append(
            # 1% collision rate at ~2M records
            Field("uid", default=nanoid.generate(size=8), unique=True)
        )
        self._metadata = Metadata(table=self.__class__.__name__, fields={f.name: f for f in self._fields})
        self.doc_id = doc_id
        super().__init__(dict({field.name: field.default for field in self._fields}, **raw_doc, **params))

    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):
        """
        Deserialize every field on the record
        """
        rec = {}
        for name, field in self._metadata.fields.items():
            rec[name] = field.deserialize(self[name], db)
        return self.__class__(rec, doc_id=self.doc_id)

    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 __repr__(self):
        return f"{self.__class__.__name__}[{self.doc_id}]: {self.items()}"


class Collection(Field):
    """
    A collection of fields that store pointers instead of dicts.
    """

    value_type = List[Record]

    def serialize(self, recs: value_type, db: TinyDB) -> List[str]:
        vals = []
        for rec in recs:
            if not rec.doc_id:
                raise PointerReferenceError(rec)
            vals.append(f"{rec._metadata.table}::{rec.uid}")
        return vals

    def deserialize(self, rec: List[str], db: TinyDB) -> Collection.value_type:
        """
        Recursively deserialize the objects in this collection
        """
        vals = []
        for member in rec:
            pt, puid = member.split("::")
            vals.append(db.table(pt).search(where("uid") == puid)[0].deserialize(db))
        return vals