vtt/src/ttfrog/schema.py

from grung.types import BackReference, Collection, Field, Record


class User(Record):
    @classmethod
    def fields(cls):
        return [*super().fields(), Field("name"), Field("email", unique=True)]


class Group(Record):
    @classmethod
    def fields(cls):
        return [*super().fields(), Field("name", unique=True), Collection("users", User)]


class Page(Record):
    @classmethod
    def fields(cls):
        return [
            *super().fields(),
            Field("uri", unique=True),
            Field("stub"),
            Field("title"),
            Field("body"),
            Collection("pages", Page),
            BackReference("parent", value_type=Page),
        ]

    def before_insert(self, db):
        super().before_insert(db)

        if not self.stub and not self.title:
            raise Exception("Must provide either a stub or a title!")
        if not self.stub:
            self.stub = self.title.title().replace(" ", "")
        if not self.title:
            self.title = self.stub

        self.uri = (self.parent.uri + "/" if self.parent and self.parent.uri != "/" else "") + self.stub

    def after_insert(self, db):
        super().after_insert(db)
        for child in self.pages:
            obj = BackReference.dereference(child, db)
            obj.uri = f"{self.uri}/{obj.stub}"
            child = db.save(obj)

    def get_child(self, obj: Record):
        for page in self.pages:
            if page.uid == obj.uid:
                return page
        return None