Coverage for nexios\exception_handler.py: 95%

1from __future__ import annotations 

2from shutil import ReadError 

3import typing 

4from nexios.exceptions import HTTPException 

5from nexios.http import Request, Response 

6from nexios.types import ExceptionHandlerType 

7from nexios.config import get_config 

8import traceback 

9from nexios.auth.exceptions import AuthenticationFailed, AuthErrorHandler 

10from nexios.exceptions import NotFoundException 

11from nexios.handlers.not_found import handle_404_error 

12from nexios import logging 

13 

14logger = logging.getLogger("nexios") 

15 

16 

17def _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 

28async 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 

61class 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 )