Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/starlette/middleware/exceptions.py: 37%

1from __future__ import annotations 

2 

3import typing 

4 

5from starlette._exception_handler import ( 

6 ExceptionHandlers, 

7 StatusHandlers, 

8 wrap_app_handling_exceptions, 

9) 

10from starlette.exceptions import HTTPException, WebSocketException 

11from starlette.requests import Request 

12from starlette.responses import PlainTextResponse, Response 

13from starlette.types import ASGIApp, Receive, Scope, Send 

14from starlette.websockets import WebSocket 

15 

16 

17class ExceptionMiddleware: 

18 def __init__( 

19 self, 

20 app: ASGIApp, 

21 handlers: typing.Mapping[typing.Any, typing.Callable[[Request, Exception], Response]] | None = None, 

22 debug: bool = False, 

23 ) -> None: 

24 self.app = app 

25 self.debug = debug # TODO: We ought to handle 404 cases if debug is set. 

26 self._status_handlers: StatusHandlers = {} 

27 self._exception_handlers: ExceptionHandlers = { 

28 HTTPException: self.http_exception, 

29 WebSocketException: self.websocket_exception, 

30 } 

31 if handlers is not None: 

32 for key, value in handlers.items(): 

33 self.add_exception_handler(key, value) 

34 

35 def add_exception_handler( 

36 self, 

37 exc_class_or_status_code: int | type[Exception], 

38 handler: typing.Callable[[Request, Exception], Response], 

39 ) -> None: 

40 if isinstance(exc_class_or_status_code, int): 

41 self._status_handlers[exc_class_or_status_code] = handler 

42 else: 

43 assert issubclass(exc_class_or_status_code, Exception) 

44 self._exception_handlers[exc_class_or_status_code] = handler 

45 

46 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: 

47 if scope["type"] not in ("http", "websocket"): 

48 await self.app(scope, receive, send) 

49 return 

50 

51 scope["starlette.exception_handlers"] = ( 

52 self._exception_handlers, 

53 self._status_handlers, 

54 ) 

55 

56 conn: Request | WebSocket 

57 if scope["type"] == "http": 

58 conn = Request(scope, receive, send) 

59 else: 

60 conn = WebSocket(scope, receive, send) 

61 

62 await wrap_app_handling_exceptions(self.app, conn)(scope, receive, send) 

63 

64 def http_exception(self, request: Request, exc: Exception) -> Response: 

65 assert isinstance(exc, HTTPException) 

66 if exc.status_code in {204, 304}: 

67 return Response(status_code=exc.status_code, headers=exc.headers) 

68 return PlainTextResponse(exc.detail, status_code=exc.status_code, headers=exc.headers) 

69 

70 async def websocket_exception(self, websocket: WebSocket, exc: Exception) -> None: 

71 assert isinstance(exc, WebSocketException) 

72 await websocket.close(code=exc.code, reason=exc.reason) # pragma: no cover