Merge pull request #25 from taskiq-python/feature/custom-errors

taskiq-python · Oct 1, 2024 · dc80597 · dc80597
2 parents 11f7437 + 152079c
commit dc80597
Showing 1 changed file with 94 additions and 47 deletions.
diff --git a/aiohttp_deps/utils.py b/aiohttp_deps/utils.py
@@ -35,6 +35,26 @@ def __init__(
         self.type_initialized = False
         self.type_cache: "Union[pydantic.TypeAdapter[Any], None]" = None
 
+    def on_validate_error(
+        self,
+        param_info: ParamInfo,
+        request: web.Request,
+        err: pydantic.ValidationError,
+    ) -> Any:
+        """Method to handle validation errors."""
+        header_name = self.alias or param_info.name
+        errors = err.errors(include_url=False)
+        for error in errors:
+            error["loc"] = (
+                "header",
+                header_name,
+            ) + error["loc"]
+            error.pop("input", None)  # type: ignore
+        raise web.HTTPBadRequest(
+            headers={"Content-Type": "application/json"},
+            text=json.dumps(errors),
+        ) from err
+
     def __call__(
         self,
         param_info: ParamInfo = Depends(),
@@ -75,17 +95,7 @@ def __call__(
         try:
             return self.type_cache.validate_python(value)
         except pydantic.ValidationError as err:
-            errors = err.errors(include_url=False)
-            for error in errors:
-                error["loc"] = (
-                    "header",
-                    header_name,
-                ) + error["loc"]
-                error.pop("input", None)  # type: ignore
-            raise web.HTTPBadRequest(
-                headers={"Content-Type": "application/json"},
-                text=json.dumps(errors),
-            ) from err
+            return self.on_validate_error(param_info, request, err)
 
 
 class Json:
@@ -100,6 +110,22 @@ def __init__(self) -> None:
         self.type_initialized = False
         self.type_cache: "Union[pydantic.TypeAdapter[Any], None]" = None
 
+    def on_validate_error(
+        self,
+        param_info: ParamInfo,
+        request: web.Request,
+        err: pydantic.ValidationError,
+    ) -> Any:
+        """Method to handle validation errors."""
+        errors = err.errors(include_url=False)
+        for error in errors:
+            error["loc"] = ("body",) + error["loc"]
+            error.pop("input", None)  # type: ignore
+        raise web.HTTPBadRequest(
+            headers={"Content-Type": "application/json"},
+            text=json.dumps(errors),
+        ) from err
+
     async def __call__(
         self,
         param_info: ParamInfo = Depends(),
@@ -135,14 +161,7 @@ async def __call__(
         try:
             return self.type_cache.validate_python(body)
         except pydantic.ValidationError as err:
-            errors = err.errors(include_url=False)
-            for error in errors:
-                error["loc"] = ("body",) + error["loc"]
-                error.pop("input", None)  # type: ignore
-            raise web.HTTPBadRequest(
-                headers={"Content-Type": "application/json"},
-                text=json.dumps(errors),
-            ) from err
+            return self.on_validate_error(param_info, request, err)
 
 
 class Query:
@@ -173,6 +192,26 @@ def __init__(
         self.type_initialized = False
         self.type_cache: "Union[pydantic.TypeAdapter[Any], None]" = None
 
+    def on_validate_error(
+        self,
+        param_info: ParamInfo,
+        request: web.Request,
+        err: pydantic.ValidationError,
+    ) -> Any:
+        """Method to handle validation errors."""
+        param_name = self.alias or param_info.name
+        errors = err.errors(include_url=False)
+        for error in errors:
+            error["loc"] = (
+                "query",
+                param_name,
+            ) + error["loc"]
+            error.pop("input", None)  # type: ignore
+        raise web.HTTPBadRequest(
+            headers={"Content-Type": "application/json"},
+            text=json.dumps(errors),
+        ) from err
+
     def __call__(
         self,
         param_info: ParamInfo = Depends(),
@@ -213,17 +252,7 @@ def __call__(
         try:
             return self.type_cache.validate_python(value)
         except pydantic.ValidationError as err:
-            errors = err.errors(include_url=False)
-            for error in errors:
-                error["loc"] = (
-                    "query",
-                    param_name,
-                ) + error["loc"]
-                error.pop("input", None)  # type: ignore
-            raise web.HTTPBadRequest(
-                headers={"Content-Type": "application/json"},
-                text=json.dumps(errors),
-            ) from err
+            return self.on_validate_error(param_info, request, err)
 
 
 class Form:
@@ -240,6 +269,22 @@ def __init__(self) -> None:
         self.type_initialized = False
         self.type_cache: "Union[pydantic.TypeAdapter[Any], None]" = None
 
+    def on_validate_error(
+        self,
+        param_info: ParamInfo,
+        request: web.Request,
+        err: pydantic.ValidationError,
+    ) -> Any:
+        """Method to handle validation errors."""
+        errors = err.errors(include_url=False)
+        for error in errors:
+            error.pop("input", None)  # type: ignore
+            error["loc"] = ("form",) + error["loc"]
+        raise web.HTTPBadRequest(
+            headers={"Content-Type": "application/json"},
+            text=json.dumps(errors),
+        ) from err
+
     async def __call__(
         self,
         param_info: ParamInfo = Depends(),
@@ -272,14 +317,7 @@ async def __call__(
         try:
             return self.type_cache.validate_python(form_data)
         except pydantic.ValidationError as err:
-            errors = err.errors(include_url=False)
-            for error in errors:
-                error.pop("input", None)  # type: ignore
-                error["loc"] = ("form",) + error["loc"]
-            raise web.HTTPBadRequest(
-                headers={"Content-Type": "application/json"},
-                text=json.dumps(errors),
-            ) from err
+            return self.on_validate_error(param_info, request, err)
 
 
 class Path:
@@ -304,6 +342,22 @@ def __init__(
         self.type_initialized = False
         self.type_cache: "Union[pydantic.TypeAdapter[Any], None]" = None
 
+    def on_validate_error(
+        self,
+        param_info: ParamInfo,
+        request: web.Request,
+        err: pydantic.ValidationError,
+    ) -> Any:
+        """Method to handle validation errors."""
+        errors = err.errors(include_url=False)
+        for error in errors:
+            error.pop("input", None)  # type: ignore
+            error["loc"] = ("path",) + error["loc"]
+        raise web.HTTPBadRequest(
+            headers={"Content-Type": "application/json"},
+            text=json.dumps(errors),
+        ) from err
+
     def __call__(
         self,
         param_info: ParamInfo = Depends(),
@@ -336,14 +390,7 @@ def __call__(
         try:
             return self.type_cache.validate_python(matched_data)
         except pydantic.ValidationError as err:
-            errors = err.errors(include_url=False)
-            for error in errors:
-                error.pop("input", None)  # type: ignore
-                error["loc"] = ("path",) + error["loc"]
-            raise web.HTTPBadRequest(
-                headers={"Content-Type": "application/json"},
-                text=json.dumps(errors),
-            ) from err
+            return self.on_validate_error(param_info, request, err)
 
 
 class ExtraOpenAPI:
@@ -375,6 +422,6 @@ def __call__(self) -> None:
         """
         This method is called when dependency is resolved.
 
-        It's empty, becuase it's used by the swagger function and
+        It's empty, becuase it's used only by the `setup_swagger` function and
         there is no actual dependency.
         """