Coverage for nexios\exception_handler.py: 95%

1 from __future__ import annotations

2 from shutil import ReadError

3 import typing

4 from nexios.exceptions import HTTPException

5 from nexios.http import Request, Response

6 from nexios.types import ExceptionHandlerType

7 from nexios.config import get_config

8 import traceback

9 from nexios.auth.exceptions import AuthenticationFailed, AuthErrorHandler

10 from nexios.exceptions import NotFoundException

11 from nexios.handlers.not_found import handle_404_error

12 from nexios import logging

13

14 logger = logging.getLogger("nexios")

15

16

17 def _lookup_exception_handler(

18 exc_handlers: typing.Dict[typing.Any[int, Exception], ExceptionHandlerType],

19 exc: Exception,

20):

21

22 for cls in type(exc).__mro__:

23 if cls in exc_handlers: # type: ignore

24 return exc_handlers[cls]

25 return None

26

27

28 async def wrap_http_exceptions(

29 request: Request,

30 response: Response,

31 call_next: typing.Callable[..., typing.Awaitable[Response]],

32 exception_handlers: typing.Dict[type[Exception], ExceptionHandlerType],

33 status_handlers: typing.Dict[int, ExceptionHandlerType],

34) -> typing.Any:

35

36 try:

37 exception_handlers, status_handlers = exception_handlers, status_handlers

38 except KeyError:

39 exception_handlers, status_handlers = {}, {}

40

41 try:

42 return await call_next()

43 except Exception as exc:

44 handler: typing.Any[ExceptionHandlerType, None] = None # type: ignore

45

46 if isinstance(exc, HTTPException):

47 handler: typing.Optional[ExceptionHandlerType] = status_handlers.get(exc.status_code) # type: ignore

48

49 if handler:

50 return await handler(request, response, exc) # type: ignore

51

52 if handler is None: # type: ignore

53 handler = _lookup_exception_handler(exception_handlers, exc)

54 if not handler:

55 error = traceback.format_exc()

56 logger.error(error)

57 raise exc

58 return await handler(request, response, exc)

59

60

61 class ExceptionMiddleware:

62 def __init__(self) -> None:

63 self.debug = (

64 get_config().debug or False

65 ) # TODO: We ought to handle 404 cases if debug is set.

66 self._status_handlers: typing.Dict[int, ExceptionHandlerType] = {}

67 self._exception_handlers: dict[

68 typing.Type[Exception], typing.Callable[..., typing.Awaitable[None]]

69 ] = {

70 HTTPException: self.http_exception,

71 AuthenticationFailed: AuthErrorHandler,

72 NotFoundException: handle_404_error,

73 }

74

75 def add_exception_handler(

76 self,

77 exc_class_or_status_code: typing.Union[int, type[Exception]],

78 handler: ExceptionHandlerType,

79 ) -> None:

80 if isinstance(exc_class_or_status_code, int):

81 self._status_handlers[exc_class_or_status_code] = handler

82 else:

83 assert issubclass(exc_class_or_status_code, Exception)

84 self._exception_handlers[exc_class_or_status_code] = handler # type:ignore

85

86 async def __call__(

87 self,

88 request: Request,

89 response: Response,

90 call_next: typing.Callable[[], typing.Awaitable[Response]],

91 ):

92

93 return await wrap_http_exceptions(

94 request=request,

95 response=response,

96 call_next=call_next,

97 exception_handlers=self._exception_handlers,

98 status_handlers=self._status_handlers,

99 )

100

101 async def http_exception(

102 self, request: Request, response: Response, exc: HTTPException

103 ) -> typing.Any:

104 assert isinstance(exc, HTTPException)

105 if exc.status_code in {204, 304}: # type:ignore

106 return response.empty(status_code=exc.status_code, headers=exc.headers)

107 return response.json(

108 exc.detail, status_code=exc.status_code, headers=exc.headers

109 )