sakhteman/core/models/base.py

import uuid
from datetime import datetime
from typing import Optional, List

import pydantic
from pydantic import Field, create_model
from pydantic.fields import ModelField, FieldInfo

from core.service.mongo_db import use_collection
from core.utils import get_time, get_uuid

MORPH_INPUT_UPDATE = ['updated_at', 'created_at', 'id', 'deleted_at', 'owner_id']
MORPH_INPUT_CREATE = ['updated_at', 'created_at', 'id', 'deleted_at', 'owner_id']


class BaseModel(pydantic.BaseModel):
    """All data models inherit from this class"""
    owner_id: Optional[uuid.UUID] = Field(description='ایجاد کننده رکورد')
    created_at: Optional[datetime] = Field(description='تاریخ ایجاد رکورد')
    updated_at: Optional[datetime] = Field(description='تاریخ آخرین بروزرسانی رکورد')

    class At:
        update: List[str]

    @pydantic.root_validator(pre=True)
    def _min_properties(cls, data):
        """At least one property is required"""
        if not data:
            raise ValueError("At least one property is required")
        return data

    def dict(self, include_nulls=False, **kwargs):
        """Override the super dict method by removing null keys from the dict, unless include_nulls=True"""
        kwargs["exclude_none"] = not include_nulls
        result = super().dict(**kwargs)
        return result

    class Config:
        extra = pydantic.Extra.ignore
        anystr_strip_whitespace = True  # strip whitespaces from strings
        underscore_attrs_are_private = False
        json_encoders = {
            uuid.UUID: lambda v: str(v),
        }

    @classmethod
    def to_partial(cls):
        return get_partial_model(cls)

    def mongo_dict(self, **kwargs):
        result = self.dict(**kwargs)
        result['_id'] = result.get('id', None)
        result.pop('id', None)
        return result

    def dict_without_id(self, **kwargs):
        result = self.dict(**kwargs)
        result.pop('id', None)
        return result

    def cu_at(self):
        self.c_at()
        self.u_at()
        return self

    def c_at(self):
        self.created_at = get_time()
        return self

    def u_at(self):
        self.updated_at = get_time()
        return self

    def init_id(self):
        self.id = get_uuid()
        return self

    def init_owner(self, is_array=False, key="owner_id"):
        from core.utils.app_context import use_user
        if is_array:
            setattr(self, key, [use_user().id if use_user() else None])
        else:
            setattr(self, key, use_user().id if use_user() else None)
        return self

    async def save(self):
        if self.DB.collection is None:
            raise ValueError("collection is not set")
        record = await use_collection(self.DB.collection).insert_one(self.u_at().mongo_dict())
        if not record.inserted_id:
            raise ValueError("Failed to save record")
        return record.inserted_id

    @classmethod
    def at(cls, op: str):
        __cols = getattr(cls.At, op)
        if __cols:
            all_fields = dict()
            for key, item in cls.__fields__.items():
                if key not in __cols:
                    if not item.required:
                        # TODO: must be correct documents
                        all_fields.setdefault(key, (item.type_, FieldInfo(
                            default=item.default,
                            description=item.field_info.description or item.__doc__,
                            default_factory=item.field_info.default_factory,
                            alias=item.field_info.alias,
                            alias_priority=item.field_info.alias_priority,
                            title=item.field_info.title,
                            exclude=item.field_info.exclude,
                            include=item.field_info.include,
                            const=item.field_info.const,
                            gt=item.field_info.gt,
                            ge=item.field_info.ge,
                            lt=item.field_info.lt,
                            le=item.field_info.le,
                            multiple_of=item.field_info.multiple_of,
                            allow_inf_nan=item.field_info.allow_inf_nan,
                            max_digits=item.field_info.max_digits,
                            decimal_places=item.field_info.decimal_places,
                            min_items=item.field_info.min_items,
                            max_items=item.field_info.max_items,
                            unique_items=item.field_info.unique_items,
                            min_length=item.field_info.min_length,
                            max_length=item.field_info.max_length,
                            allow_mutation=item.field_info.allow_mutation,
                            repr=item.field_info.repr,
                            regex=item.field_info.regex,
                            discriminator=item.field_info.discriminator,
                            extra=item.field_info.extra,
                        )))
                    else:
                        all_fields.setdefault(key, (item.type_, ...))
            DynamicModel = create_model(cls.__name__, **all_fields)
            return DynamicModel
        return cls


class BaseSoftDelete(pydantic.BaseModel):
    deleted_at: Optional[datetime] = Field(title='زمان حذف رکورد',
                                           description='زمانی که این فیلد در داده وجود داشته باشد نشان دهنده این است که این رکورد حذف نرم شده است.')


class BaseCUAt(pydantic.BaseModel):
    created_at: Optional[datetime]
    updated_at: Optional[datetime]


DELETED_AT_QUERY = {'deleted_at': {'$exists': False}}

ORDER_BY_CREATED_AT = {'key_or_list': 'created_at', 'direction': -1}


def get_partial_model(model: BaseModel) -> BaseModel:
    """
    Return a model similar to the original, but where all fields are Optional.

    Note: this minimal implementation means that many Pydantic features will be discarded
    (such as alias). This is not by design and could stand to improve.
    """
    new_fields = {name: (Optional[model.__annotations__.get(name)], model.__fields__[
        name].default or None) for name in model.__fields__}

    return pydantic.create_model(
        "PartialModel",
        **new_fields,
        __base__=BaseModel)