DEV Community

Cover image for LGTM Devlog 29: ORM for Firestore and __init__subclass__ dunders and metaclasses
Yuan Gao
Yuan Gao

Posted on

LGTM Devlog 29: ORM for Firestore and __init__subclass__ dunders and metaclasses

Uh oh, I refactored again. I wasn't quite happy with how messy it was, how each Game, User, and Quest entity was handling their data, there was a lot of duplicate code when it came to communicating with the database, and when that happened.

For example, the game object prior to the latest change looks like this:

class Game:
    @classmethod
    def from_user(cls, user: User) -> Union[Game, NoGameType]:
        """ Create a game from a user """
        key = cls.make_key(user)
        game = cls(key)
        game.user = user

        docs = db.collection("game").where("user_key", "==", user.key).stream()
        for _ in docs:
            return game
        return NoGame

    @classmethod
    def new_from_fork(cls, user: User, fork_url: str) -> Game:
        """ Save game with fork """

        if not fork_url:
            raise ValueError("Fork can't be blank")

        key = cls.make_key(user)
        game = cls(key)
        game.user = user

        game_doc = db.collection("game").document(game.key).get()
        if game_doc.exists:
            game_doc.reference.set(
                {
                    "fork_url": fork_url,
                    "user_uid": user.uid,
                    "user_key": user.key,
                },
                merge=True,
            )

        else:
            game_doc.reference.set(
                {
                    "fork_url": fork_url,
                    "user_uid": user.uid,
                    "user_key": user.key,
                    "joined": firestore.SERVER_TIMESTAMP,
                }
            )

        return game

    @staticmethod
    def make_key(user: User) -> str:
        """ Game's key ARE user key due to 1:1 relationship """
        return user.key

    key: str
    user: Union[User, NoUserType]

    def __init__(self, key: str):
        self.key = key
        self.user = NoUser

    def assign_to_uid(self, uid: str) -> None:
        """ Assign a user to this game """
        doc = db.collection("game").document(self.key).get()
        if doc.exists:
            doc.reference.set({"user_uid": uid}, merge=True)

    def start_first_quest(self) -> None:
        """ Create starting quest if not exist """
        QuestClass = Quest.get_first_quest()
        quest = QuestClass.from_game(self)
        quest.execute_stages(TickType.FULL)
        quest.save()

    def __repr__(self):
        return f"{self.__class__.__name__}(key={self.key})"
Enter fullscreen mode Exit fullscreen mode

We have two functions for creating Game objects: new_from_fork() and from_user(). There's also the method assign_to_uid() which writes a single property to the database.

This seems ok, but there's a lot of duplicate code in the User object, and Quest objects.

So I decided to warp a lot of this duplicate function into an ORM base class so that User, Quest and Game can all inherit these functions that will save and restore themselves from the database.

After the change, the new Game object looks like this:

class Game(Orm, collection="game", parent_orm=User):
    data: GameData
    storage_model = GameData

    @classmethod
    def from_user(cls, user: User) -> Game:
        key = cls.make_key(user)
        game = cls(key)
        game.parent_key = user.key
        return game

    @staticmethod
    def make_key(user: User) -> str:
        """ Game's key ARE user key due to 1:1 relationship """
        return user.key

    def set_fork_url(self, fork_url: str) -> None:
        self.data.fork_url = fork_url
Enter fullscreen mode Exit fullscreen mode

Much cleaner! The first thing you'll notice is the line:

class Game(Orm, collection="game", parent_orm=User):
Enter fullscreen mode Exit fullscreen mode

This is a neat feature in Python allowing customisation of metaclasses and is described in PEP 487. It's powered by the __init_subclass__ dunder, which runs when the class is initialised (compared to __init__ when the instance is initialised), which in this case looks like this (from the Orm base class):

    def __init_subclass__(
        cls, collection: str, parent_orm: Union[Type[Orm], NoParentType] = NoParent
    ):
        """ Set collection and parent """
        cls.collection = collection
        cls.parent_orm = parent_orm
        cls.col_ref = db.collection(collection)
Enter fullscreen mode Exit fullscreen mode

So, what is happening here is when Game class is initialized, it is setting the collection class variable, and the parent_orm class variable. While we could have used __init__ to set the collection and parent_orm, this method provides a nice separation of concerns between the Game class, whose class properties relate to the collection of games; and a specific instance of the Game class which relate to a specific entry in this collection. Setting col_ref in __init_subclass__ allows us to do Game.col_ref without having to instantiate a Game object (which we shouldn't need to do when referencing all Games).

The other changes is removing new_from_fork() and replacing with a short set_fork_url() setter. The reason for this is I realised both new_from_fork() and from_user() require a user argument, so there's a duplication here. assign_to_uid() can now use the Orm base class's parent_key property; and the quest stuff has been moved to a new QuestPage object instead, where it's more relevant.

The ORM

The Underlying ORM base class's job is to map data between Firestore objects and Python objects. It has all the shared code that Game, User and new QuestPage object needs. The top half of it looks like this:


class Orm(ABC):
    """ ORM base class links stuff together """

    @property
    @abstractmethod
    def storage_model(cls) -> Type[BaseModel]:
        """ Storage model """
        return NotImplemented

    collection: ClassVar[str]
    parent_orm: ClassVar[Union[Type[Orm], NoParentType]]
    col_ref: CollectionReference

    def __init_subclass__(
        cls, collection: str, parent_orm: Union[Type[Orm], NoParentType] = NoParent
    ):
        """ Set collection and parent """
        cls.collection = collection
        cls.parent_orm = parent_orm
        cls.col_ref = db.collection(collection)

    key: Union[str, NoKeyType]
    parent_key: Union[str, NoKeyType]
    data: BaseModel

    def __init__(self, key: Union[str, NoKeyType] = NoKey):
        self.key = key
        self.data = self.storage_model()
        self.parent_key = NoKey

    @property
    def parent(self) -> Union[Orm, OrmNotFoundType]:
        if self.parent_orm is not NoParent and self.parent_key is not NoKey:
            return self.parent_orm(self.parent_key)

        return OrmNotFound

    ...
Enter fullscreen mode Exit fullscreen mode

With the bottom half related to the mechanics of saving and loading data from the database.

I've described the __init_subclass__ already, the rest of this ABC enforces that concrete implementations should supply their own storage_model, a Pydantic model which the base class's methods will use for creating and restoring data firestore. It also has a property called parent which will return an instance of the parent object linked by a parent_key property. For example if we had an instantiated Game object, we could fetch it's user by doing:

game = Game(key)
game.load()

user = game.parent
Enter fullscreen mode Exit fullscreen mode

Because Game objects know that their parent ORM object is User, and Game objects also store a parent_key property. This is enough detail to return a User object.

QuestPages

I've done some big refactoring to the former Quest object. It's now been split into two: Quest which contains everything related to running the Quest, and serves as a base class to concrete implementations of individual Quests; and QuestPage which is subclassed from Orm and relates only to saving and loading quests from the database. I imagine QuestPage as a page in a quest log, which contains the "save data" relating to a Quest which is the actual implementation of the quest.

There's a bit of composition going on, as QuestPage has one instance of quest as a property.


class QuestPage(Orm, collection="quest", parent_orm=Game):
    data: QuestData
    storage_model = QuestData
    quest: Quest

    @staticmethod
    def make_key(game: Game, quest_name: str) -> str:
        if not quest_name:
            raise ValueError("quest_name must be valid")
        return f"{game.key}:{quest_name}"

    @classmethod
    def from_game_get_first_quest(cls, game: Game) -> QuestPage:
        return cls.from_game_get_quest(game, FIRST_QUEST_NAME)

    @classmethod
    def from_game_get_quest(cls, game: Game, quest_name: str) -> QuestPage:
        key = cls.make_key(game, quest_name)
        quest = cls(key, quest_name)
        return quest

    @classmethod
    def iterate_all(cls) -> Generator[QuestPage, None, None]:
        """ Iterate over all quests, the generator yields loaded quest_pages """
        docs = cls.col_ref.where("complete", "!=", True).stream()
        for doc in docs:
            data = cls.storage_model.parse_obj(doc.to_dict())
            quest_page = cls(doc.id, data.quest_name)
            quest_page.data = data
            quest_page.quest.load_raw(data.version, data.serialized_data)
            yield quest_page

    def __init__(self, key: str, quest_name):
        super().__init__(key)
        self.quest = Quest.from_name(quest_name, self)
        self.data.quest_name = quest_name

    def load(self) -> None:
        """ Additionally parse the quest storage """
        super().load()
        if isinstance(self.quest, Quest):
            self.quest.load_raw(self.data.version, self.data.serialized_data)

    def save(self) -> None:
        """ Additionally parse out the quest storage """
        if isinstance(self.quest, Quest):
            self.data.serialized_data = self.quest.save_raw()
            self.data.version = str(self.quest.version)
        super().save()
    ...
Enter fullscreen mode Exit fullscreen mode

This extends the __init__(), save(), and load() of the ORM base class in order to also save and load the quest data model distinct from the quest_page data model. The reason for this split is so that the quests themselves have their own variable space for quest-related data, while all the metadata - the version number, the completed quest list, etc. are store separately.


Together with the other changes, this refactor, which still hasn't added any new functionality, has cleaned up the code once again. Hopefully i don't have to do too many more of these

Top comments (0)