allow except mapping in class

restit/common.py

@@ -1,5 +1,6 @@
 import os
 import sys
+from functools import lru_cache
 from html import escape
 from typing import List
 
@@ -50,7 +51,15 @@ def guess_text_content_subtype_string(content: str) -> str:
     return "text/plain"
 
 
+@lru_cache()
 def get_response_status_parameters_for_method(method_object: object) -> List[ResponseStatusParameter]:
     response_status_parameters = getattr(method_object.__self__, "__response_status_parameters__", [])
     response_status_parameters.extend(getattr(method_object, "__response_status_parameters__", []))
     return response_status_parameters
+
+
+@lru_cache()
+def get_exception_mapping_for_method(method_object: object) -> dict:
+    exception_mapping: dict = getattr(method_object.__self__, "__exception_mapping__", {})
+    exception_mapping.update(getattr(method_object, "__exception_mapping__", {}))
+    return exception_mapping

restit/decorator/exception_mapping_decorator.py

@@ -7,27 +7,8 @@
 
 
 def exception_mapping(mapping: Dict[Type[Exception], Union[Tuple[Type[HttpError], str], Type[HttpError]]]):
-    def decorator(func):
-        def wrapper(*args, **kwargs):
-            try:
-                return func(*args, **kwargs)
-            except Exception as exception:
-                for source_exception_class, target_exception_tuple_or_class in mapping.items():
-                    if isinstance(exception, source_exception_class):
-                        if isinstance(target_exception_tuple_or_class, tuple):
-                            LOGGER.debug(
-                                "Mapping exception class %s to %s with description: %s",
-                                type(exception), target_exception_tuple_or_class[0], target_exception_tuple_or_class[1]
-                            )
-                            raise target_exception_tuple_or_class[0](target_exception_tuple_or_class[1])
-                        else:
-                            LOGGER.debug(
-                                "Mapping exception class %s to %s", type(exception), target_exception_tuple_or_class
-                            )
-                            raise target_exception_tuple_or_class(str(exception))
-
-                raise exception
-
-        return wrapper
+    def decorator(func_or_class):
+        setattr(func_or_class, "__exception_mapping__", mapping)
+        return func_or_class
 
     return decorator

restit/resource.py

@@ -2,13 +2,13 @@
 import inspect
 import logging
 import re
-from typing import Tuple, AnyStr, Dict, Union, List
+from typing import Tuple, AnyStr, Dict, Union, List, Callable
 
 from marshmallow import ValidationError
 
 from restit._path_parameter import PathParameter
 from restit._response import Response
-from restit.common import get_response_status_parameters_for_method
+from restit.common import get_response_status_parameters_for_method, get_exception_mapping_for_method
 from restit.exception import MethodNotAllowed
 from restit.exception.client_errors_4xx import BadRequest
 from restit.internal.query_parameter import QueryParameter
@@ -94,7 +94,7 @@ def handle_request(self, request_method: str, request: Request, path_params: Dic
         request._path_params = self._collect_and_convert_path_parameters(path_params)
         self._process_query_parameters(method_object, request)
         request = self._validate_request_body(method_object, request)
-        response: Response = method_object(request)
+        response: Response = self._execute_request_with_exception_mapping(method_object, request)
         if not isinstance(response, Response):
             raise Resource.NoResponseReturnException(
                 f"Resource method {method_object} does not return a response object"
@@ -105,6 +105,28 @@ def handle_request(self, request_method: str, request: Request, path_params: Dic
         )
         return response
 
+    @staticmethod
+    def _execute_request_with_exception_mapping(method_object: Callable, request: Request) -> Response:
+        exception_mapping = get_exception_mapping_for_method(method_object)
+        try:
+            return method_object(request)
+        except Exception as exception:
+            for source_exception_class, target_exception_tuple_or_class in exception_mapping.items():
+                if isinstance(exception, source_exception_class):
+                    if isinstance(target_exception_tuple_or_class, tuple):
+                        LOGGER.debug(
+                            "Mapping exception class %s to %s with description: %s",
+                            type(exception), target_exception_tuple_or_class[0], target_exception_tuple_or_class[1]
+                        )
+                        raise target_exception_tuple_or_class[0](target_exception_tuple_or_class[1])
+                    else:
+                        LOGGER.debug(
+                            "Mapping exception class %s to %s", type(exception), target_exception_tuple_or_class
+                        )
+                        raise target_exception_tuple_or_class(str(exception))
+
+            raise exception
+
     @staticmethod
     def _find_response_schema_by_status(status: int, method_object: object) -> Union[None, ResponseStatusParameter]:
         response_status_parameters = get_response_status_parameters_for_method(method_object)

test/restit/exception_mapping_test.py

@@ -0,0 +1,37 @@
+import unittest
+
+from restit import Resource, Request, Response, RestItTestApp, RestItApp
+from restit.decorator import path, exception_mapping
+from restit.exception import BadRequest, NotFound
+
+
+class MyException1(Exception):
+    pass
+
+
+@path("/")
+@exception_mapping({MyException1: BadRequest})
+class MyResource(Resource):
+    def get(self, request: Request) -> Response:
+        raise MyException1("Hello")
+
+    @exception_mapping({MyException1: (NotFound, "Miau")})
+    def post(self, request: Request) -> Response:
+        raise MyException1()
+
+
+class ExceptionMappingTestTestCase(unittest.TestCase):
+    def setUp(self) -> None:
+        self.restit_test_app = RestItTestApp(RestItApp(resources=[MyResource()], raise_exceptions=True))
+
+    def test_map_exception(self):
+        with self.assertRaises(BadRequest) as exception:
+            self.restit_test_app.get("/")
+
+        self.assertEqual("Hello", str(exception.exception))
+
+    def test_override(self):
+        with self.assertRaises(NotFound) as exception:
+            self.restit_test_app.post("/")
+
+        self.assertEqual("Miau", str(exception.exception))