feat: add convenience hooks for to_model operations (#347)

* feat: add convenience hooks on `to_model` * fix: swap to walrus
litestar-org · Jan 17, 2025 · d7ce1c2 · d7ce1c2
1 parent 20e0ab6
commit d7ce1c2
Show file tree

Hide file tree

Showing 2 changed files with 104 additions and 0 deletions.
diff --git a/advanced_alchemy/service/_async.py b/advanced_alchemy/service/_async.py
@@ -333,6 +333,50 @@ async def get_one_or_none(
             ),
         )
 
+    async def to_model_on_create(self, data: ModelDictT[ModelT]) -> ModelDictT[ModelT]:
+        """Convenience method to allow for custom behavior on create.
+
+        Args:
+            data: The data to be converted to a model.
+
+        Returns:
+            The data to be converted to a model.
+        """
+        return data
+
+    async def to_model_on_update(self, data: ModelDictT[ModelT]) -> ModelDictT[ModelT]:
+        """Convenience method to allow for custom behavior on update.
+
+        Args:
+            data: The data to be converted to a model.
+
+        Returns:
+            The data to be converted to a model.
+        """
+        return data
+
+    async def to_model_on_delete(self, data: ModelDictT[ModelT]) -> ModelDictT[ModelT]:
+        """Convenience method to allow for custom behavior on delete.
+
+        Args:
+            data: The data to be converted to a model.
+
+        Returns:
+            The data to be converted to a model.
+        """
+        return data
+
+    async def to_model_on_upsert(self, data: ModelDictT[ModelT]) -> ModelDictT[ModelT]:
+        """Convenience method to allow for custom behavior on upsert.
+
+        Args:
+            data: The data to be converted to a model.
+
+        Returns:
+            The data to be converted to a model.
+        """
+        return data
+
     async def to_model(
         self,
         data: ModelDictT[ModelT],
@@ -346,6 +390,14 @@ async def to_model(
         Returns:
             Representation of created instances.
         """
+        operation_map = {
+            "create": self.to_model_on_create,
+            "update": self.to_model_on_update,
+            "delete": self.to_model_on_delete,
+            "upsert": self.to_model_on_upsert,
+        }
+        if operation and (op := operation_map.get(operation)):
+            data = await op(data)
         if is_dict(data):
             return model_from_dict(model=self.model_type, **data)
         if is_pydantic_model(data):

diff --git a/advanced_alchemy/service/_sync.py b/advanced_alchemy/service/_sync.py
@@ -347,6 +347,50 @@ def get_one_or_none(
             ),
         )
 
+    def to_model_on_create(self, data: ModelDictT[ModelT]) -> ModelDictT[ModelT]:
+        """Convenience method to allow for custom behavior on create.
+
+        Args:
+            data: The data to be converted to a model.
+
+        Returns:
+            The data to be converted to a model.
+        """
+        return data
+
+    def to_model_on_update(self, data: ModelDictT[ModelT]) -> ModelDictT[ModelT]:
+        """Convenience method to allow for custom behavior on update.
+
+        Args:
+            data: The data to be converted to a model.
+
+        Returns:
+            The data to be converted to a model.
+        """
+        return data
+
+    def to_model_on_delete(self, data: ModelDictT[ModelT]) -> ModelDictT[ModelT]:
+        """Convenience method to allow for custom behavior on delete.
+
+        Args:
+            data: The data to be converted to a model.
+
+        Returns:
+            The data to be converted to a model.
+        """
+        return data
+
+    def to_model_on_upsert(self, data: ModelDictT[ModelT]) -> ModelDictT[ModelT]:
+        """Convenience method to allow for custom behavior on upsert.
+
+        Args:
+            data: The data to be converted to a model.
+
+        Returns:
+            The data to be converted to a model.
+        """
+        return data
+
     def to_model(
         self,
         data: ModelDictT[ModelT],
@@ -360,6 +404,14 @@ def to_model(
         Returns:
             Representation of created instances.
         """
+        operation_map = {
+            "create": self.to_model_on_create,
+            "update": self.to_model_on_update,
+            "delete": self.to_model_on_delete,
+            "upsert": self.to_model_on_upsert,
+        }
+        if operation and (op := operation_map.get(operation)):
+            data = op(data)
         if is_dict(data):
             return model_from_dict(model=self.model_type, **data)
         if is_pydantic_model(data):