from datetime import date, datetime from typing import List, Literal, Optional from pydantic import Field, BaseModel, validator from beanie import Indexed from ..core.models import CrudDocument class EntityType(BaseModel): @property def label(self) -> str: return self.title class Individual(EntityType): type: Literal['individual'] = 'individual' firstname: Indexed(str) middlename: Indexed(str) = "" lastname: Indexed(str) surnames: List[Indexed(str)] = [] day_of_birth: date place_of_birth: str = "" @property def label(self) -> str: if len(self.surnames) > 0: return '{} "{}" {}'.format(self.firstname, self.surnames[0], self.lastname) return '{} {}'.format(self.firstname, self.lastname) class Employee(BaseModel): role: Indexed(str) entity_id: str = Field(foreignKey={ "reference": { "resource": "entity", "schema": "Entity", "condition": "entity_data.type=individual" } }) class Corporation(EntityType): type: Literal['corporation'] = 'corporation' title: Indexed(str) activity: Indexed(str) employees: List[Employee] = Field(default=[]) class Institution(EntityType): type: Literal['institution'] = 'institution' title: Indexed(str) activity: Indexed(str) employees: List[Employee] = Field(default=[]) class Entity(CrudDocument): entity_data: Individual | Corporation | Institution = Field(..., discriminator='type') label: str = None address: Optional[str] = "" @validator("label", always=True) def generate_label(cls, v, values, **kwargs): if 'entity_data' not in values: return v return values['entity_data'].label class Settings(CrudDocument.Settings): fulltext_search = ['label'] bson_encoders = { date: lambda dt: dt if hasattr(dt, 'hour') else datetime(year=dt.year, month=dt.month, day=dt.day, hour=0, minute=0, second=0) }