accessing child class attributes from parent static methods in pydantic models-CodePudding

Given a pydantic BaseModel class defined as follows:

from typing import List, Optional
from uuid import uuid4

from pydantic import BaseModel, Field
from server.database import get_db


class Campaign(BaseModel):
    id: Optional[str] = Field(default_factory=lambda: str(uuid4()))
    name: str

    @staticmethod
    async def all() -> List:
        ret = get_db()['campaigns'].find()
        return [Campaign(**i) async for i in ret]

    @staticmethod
    async def get(id):
        ret = await get_db()['campaigns'].find_one({'id': id})
        return Campaign(**ret)

    async def save(self):
        await get_db()['campaigns'].insert_one(self.dict())

    async def update(self, **kwargs):
        await get_db()['campaigns'].update_one(
            {'id': self.id},
            {'$set': kwargs},
        )

As you can see from the above model defination that the functions [all, get, save, update] can be universal for multiple CRUD models like as follows:

class Record(BaseModel):
    ...
    @staticmethod
    async def all() -> List:
        ret = get_db()['records'].find()
        return [Record(**i) async for i in ret]

    async def save(self):
        pass

class Fragment(BaseModel):
    ...
    @staticmethod
    async def all() -> List:
        ret = get_db()['fragments'].find()
        return [Fragment(**i) async for i in ret]

    async def save(self):
        pass

Across all these methods of different classes the only change is the name of the MongoDB collection. Right now, I am copy pasting this code multiple times for different models.

Is there any generic solution so that I can encapsulate the definition of these methods under a single base class and hopefully follow the DRY principle (as opposed to blatantly breaking it right now).

One such code configuration can be:

class BaseDBModel(BaseModel):
    async def save(self):
        await get_db()[self.__class__.Meta.collection_name].insert_one(self.dict())

    async def update(self, **kwargs):
        await get_db()[self.__class__.Meta.collection_name].update_one(
            {'id': self.id},
            {'$set': kwargs},
        )

class Campaign(BaseDBModel):
    id: Optional[str] = Field(default_factory=lambda: str(uuid4()))
    name: str

    class Meta:
        collection_name = 'campaigns'

But, this code will only work for instance methods (i.e. save, update) and would not work for static methods like all and get. I am unsure as to how a parent class staticmethod will go about accessing the Meta class of the children in python?

Any suggestions or improvements on the above approach would be appreciated.

CodePudding user response：

The static methods should be class methods instead, so that the class invoking the method is passed as the first argument. For example,

class BaseDBModel(BaseModel):

    @classmethod
    async def all(cls) -> List:
        ret = get_db()[cls.Meta.collection_name].find()
        return [class(**i) async for i in ret]

    @classmethod
    async def get(id):
        ret = await get_db()[cls.Meta.collection_name]({'id': id})
        return class(**ret)

    async def save(self):
        await get_db()[self.Meta.collection_name].insert_one(self.dict())

    async def update(self, **kwargs):
        await get_db()[self.Meta.collection_name].update_one(
            {'id': self.id},
            {'$set': kwargs},
        )


class Campaign(BaseDBModel):
    id: Optional[str] = Field(default_factory=lambda: str(uuid4()))
    name: str
    
    class Meta:
        collection_name = 'campaigns'


class Record(Campaign):
    class Meta:
        collection_name = 'records'


class Fragment(Campaign):
    class Meta:
        collection_name = 'fragments'

(It's not clear if Record and Fragment should inherit from Campaign or BaseDBModel directly.)