Adding support for nested types, array and dates!

back/app/entity/models.py

@@ -1,8 +1,9 @@
 from enum import Enum
-from datetime import datetime
-from pydantic import Field
+from datetime import datetime, date
+from typing import List, Literal
 
-from beanie import Document
+from pydantic import Field, BaseModel
+from beanie import Document, Link
 
 
 class EntityType(str, Enum):
@@ -11,13 +12,40 @@ class EntityType(str, Enum):
     institution = 'institution'
 
 
+class Individual(BaseModel):
+    type: Literal['individual'] = 'individual'
+    firstname: str
+    middlenames: List[str] = Field(default=[])
+    lastname: str
+    surnames: List[str] = Field(default=[])
+    day_of_birth: date
+    job: str
+    employer: str
+
+
+class Employee(BaseModel):
+    entity_id: str
+    role: str
+
+
+class Corporation(BaseModel):
+    type: Literal['corporation'] = 'corporation'
+    title: str
+    activity: str
+    employees: List[Employee] = Field(default=[])
+
+
 class Entity(Document):
     _id: str
-    type: EntityType
     name: str
     address: str
+    entity_data: Individual | Corporation = Field(..., discriminator='type')
+
     created_at: datetime = Field(default=datetime.utcnow(), nullable=False)
     updated_at: datetime = Field(default_factory=datetime.utcnow, nullable=False)
-    #
-    # class Settings:
-    #     name = "entities"
+
+    class Settings:
+        bson_encoders = {
+            date: lambda dt: datetime(year=dt.year, month=dt.month, day=dt.day, hour=0, minute=0, second=0) \
+                                    if not hasattr(dt, 'hour') else dt
+        }