Files
cht-lawfirm/back/app/entity/models.py
2023-01-26 16:53:14 +01:00

83 lines
2.3 KiB
Python

from enum import Enum
from datetime import datetime, date
from typing import List, Literal, Optional
from pymongo import TEXT, IndexModel
from pydantic import Field, BaseModel, validator
from beanie import Document, Indexed
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
@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",
"displayName": "_id",
"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(Document):
_id: str
entity_data: Individual | Corporation | Institution = Field(..., discriminator='type')
address: Optional[str] = ""
label: str = None
@validator("label", always=True)
def generate_label(cls, v, values, **kwargs):
if 'entity_data' not in values:
return v
return values['entity_data'].label
created_at: datetime = Field(default=datetime.utcnow(), nullable=False)
updated_at: datetime = Field(default_factory=datetime.utcnow, nullable=False)
class Settings:
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)
}
fulltext_search = ['label']