Adding filter on Entity type

api/rpk-api/firm/entity/models.py

@@ -1,10 +1,12 @@
 from datetime import date
+from enum import Enum
 from typing import List, Literal, Optional
 
+from fastapi_filter import FilterDepends, with_prefix
 from pydantic import Field, BaseModel, ConfigDict
 from beanie import PydanticObjectId
 
-from firm.core.models import CrudDocument, ForeignKey
+from firm.core.models import CrudDocument, ForeignKey, CrudDocumentConfig
 from firm.core.filter import Filter, FilterSchema
 
 
@@ -13,6 +15,10 @@ class EntityType(BaseModel):
     def label(self) -> str:
         return self.title
 
+class EntityTypeEnum(str, Enum):
+    individual = 'individual'
+    corporation = 'corporation'
+    institution = 'institution'
 
 class Individual(EntityType):
     model_config = ConfigDict(title='Particulier')
@@ -63,6 +69,9 @@ class Entity(CrudDocument):
     Fiche d'un client
     """
     model_config = ConfigDict(title='Client')
+    document_config = CrudDocumentConfig(
+        indexes=["entity_data.type"],
+    )
 
     entity_data: Individual | Corporation | Institution = Field(..., discriminator='type')
     address: str = Field(default="", title='Adresse')
@@ -73,7 +82,16 @@ class Entity(CrudDocument):
         return self.entity_data.label
 
 
+class EntityDataFilter(Filter):
+    type__in: Optional[list[EntityTypeEnum]] = None
+
+    class Constants(Filter.Constants):
+        model = EntityType
+
+
 class EntityFilters(FilterSchema):
+    entity_data: Optional[EntityDataFilter] = FilterDepends(with_prefix("entity_data", EntityDataFilter))
+
     class Constants(Filter.Constants):
         model = Entity
         search_model_fields = ["label"]