encode · kristjanvalur · Feb 26, 2023 · Feb 26, 2023 · Feb 26, 2023
diff --git a/starlette/applications.py b/starlette/applications.py
@@ -5,7 +5,7 @@
 from starlette.middleware import Middleware
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.middleware.errors import ServerErrorMiddleware
-from starlette.middleware.exceptions import ExceptionMiddleware
+from starlette.middleware.exceptions import ExceptionHandler, ExceptionMiddleware
 from starlette.requests import Request
 from starlette.responses import Response
 from starlette.routing import BaseRoute, Router
@@ -47,9 +47,12 @@ def __init__(
         exception_handlers: typing.Optional[
             typing.Mapping[
                 typing.Any,
-                typing.Callable[
-                    [Request, Exception],
-                    typing.Union[Response, typing.Awaitable[Response]],
+                typing.Union[
+                    ExceptionHandler,
+                    typing.Callable[
+                        [Request, Exception],
+                        typing.Union[Response, typing.Awaitable[Response]],
+                    ],
                 ],
             ]
         ] = None,
@@ -80,7 +83,14 @@ def build_middleware_stack(self) -> ASGIApp:
         debug = self.debug
         error_handler = None
         exception_handlers: typing.Dict[
-            typing.Any, typing.Callable[[Request, Exception], Response]
+            typing.Any,
+            typing.Union[
+                ExceptionHandler,
+                typing.Callable[
+                    [Request, Exception],
+                    typing.Union[Response, typing.Awaitable[Response]],
+                ],
+            ],
         ] = {}
 
         for key, value in self.exception_handlers.items():

diff --git a/starlette/middleware/errors.py b/starlette/middleware/errors.py
@@ -5,6 +5,7 @@
 
 from starlette._utils import is_async_callable
 from starlette.concurrency import run_in_threadpool
+from starlette.middleware.exceptions import ExceptionHandler
 from starlette.requests import Request
 from starlette.responses import HTMLResponse, PlainTextResponse, Response
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
@@ -137,11 +138,17 @@ class ServerErrorMiddleware:
     def __init__(
         self,
         app: ASGIApp,
-        handler: typing.Optional[typing.Callable] = None,
+        handler: typing.Optional[
+            typing.Union[
+                ExceptionHandler, typing.Callable[[typing.Any, typing.Any], typing.Any]
+            ]
+        ] = None,
         debug: bool = False,
     ) -> None:
         self.app = app
-        self.handler = handler
+        self.handler = (
+            handler.handler if isinstance(handler, ExceptionHandler) else handler
+        )
         self.debug = debug
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:

diff --git a/starlette/middleware/exceptions.py b/starlette/middleware/exceptions.py
@@ -8,13 +8,43 @@
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 from starlette.websockets import WebSocket
 
+_ET = typing.TypeVar("_ET", bound=Exception)
+
+
+class ExceptionHandler:
+    @typing.overload
+    def __init__(
+        self,
+        handler: typing.Callable[
+            [Request, _ET], typing.Union[Response, typing.Awaitable[Response]]
+        ],
+    ) -> None:
+        ...
+
+    @typing.overload
+    def __init__(
+        self,
+        handler: typing.Callable[
+            [WebSocket, _ET], typing.Union[None, typing.Awaitable[None]]
+        ],
+    ) -> None:
+        ...
+
+    def __init__(self, handler: typing.Callable[[typing.Any, _ET], typing.Any]) -> None:
+        self.handler: typing.Callable[[typing.Any, typing.Any], typing.Any] = handler
+
 
 class ExceptionMiddleware:
     def __init__(
         self,
         app: ASGIApp,
         handlers: typing.Optional[
-            typing.Mapping[typing.Any, typing.Callable[[Request, Exception], Response]]
+            typing.Mapping[
+                typing.Any,
+                typing.Union[
+                    ExceptionHandler, typing.Callable[[Request, Exception], Response]
+                ],
+            ]
         ] = None,
         debug: bool = False,
     ) -> None:
@@ -31,11 +61,31 @@ def __init__(
             for key, value in handlers.items():
                 self.add_exception_handler(key, value)
 
+    @typing.overload
+    def add_exception_handler(
+        self,
+        exc_class_or_status_code: typing.Union[int, typing.Type[Exception]],
+        handler: ExceptionHandler,
+    ) -> None:
+        ...
+
+    @typing.overload
     def add_exception_handler(
         self,
         exc_class_or_status_code: typing.Union[int, typing.Type[Exception]],
         handler: typing.Callable[[Request, Exception], Response],
     ) -> None:
+        ...
+
+    def add_exception_handler(
+        self,
+        exc_class_or_status_code: typing.Union[int, typing.Type[Exception]],
+        handler: typing.Union[
+            ExceptionHandler, typing.Callable[[typing.Any, typing.Any], typing.Any]
+        ],
+    ) -> None:
+        if isinstance(handler, ExceptionHandler):
+            handler = handler.handler
         if isinstance(exc_class_or_status_code, int):
             self._status_handlers[exc_class_or_status_code] = handler
         else:
@@ -44,7 +94,7 @@ def add_exception_handler(
 
     def _lookup_exception_handler(
         self, exc: Exception
-    ) -> typing.Optional[typing.Callable]:
+    ) -> typing.Optional[typing.Callable[[typing.Any, typing.Any], typing.Any]]:
         for cls in type(exc).__mro__:
             if cls in self._exception_handlers:
                 return self._exception_handlers[cls]

diff --git a/tests 8000 /test_applications.py b/tests 8000 /test_applications.py
@@ -7,27 +7,28 @@
 import pytest
 
 from starlette import status
-from starlette.applications import Starlette
+from starlette.applications import ExceptionHandler, Starlette
 from starlette.endpoints import HTTPEndpoint
 from starlette.exceptions import HTTPException, WebSocketException
 from starlette.middleware import Middleware
 from starlette.middleware.trustedhost import TrustedHostMiddleware
-from starlette.responses import JSONResponse, PlainTextResponse
+from starlette.requests import Request
+from starlette.responses import JSONResponse, PlainTextResponse, Response
 from starlette.routing import Host, Mount, Route, Router, WebSocketRoute
 from starlette.staticfiles import StaticFiles
 from starlette.types import ASGIApp
 from starlette.websockets import WebSocket
 
 
-async def error_500(request, exc):
+async def error_500(request: Request, exc: Exception) -> Response:
     return JSONResponse({"detail": "Server Error"}, status_code=500)
 
 
-async def method_not_allowed(request, exc):
+async def method_not_allowed(request: Request, exc: Exception) -> Response:
     return JSONResponse({"detail": "Custom message"}, status_code=405)
 
 
-async def http_exception(request, exc):
+async def http_exception(request: Request, exc: HTTPException) -> Response:
     return JSONResponse({"detail": exc.detail}, status_code=exc.status_code)
 
 
@@ -67,7 +68,7 @@ async def websocket_endpoint(session):
     await session.close()
 
 
-async def websocket_raise_websocket(websocket: WebSocket):
+async def websocket_raise_websocket(websocket: WebSocket) -> None:
     await websocket.accept()
     raise WebSocketException(code=status.WS_1003_UNSUPPORTED_DATA)
 
@@ -76,12 +77,12 @@ class CustomWSException(Exception):
     pass
 
 
-async def websocket_raise_custom(websocket: WebSocket):
+async def websocket_raise_custom(websocket: WebSocket) -> None:
     await websocket.accept()
     raise CustomWSException()
 
 
-def custom_ws_exception_handler(websocket: WebSocket, exc: CustomWSException):
+def custom_ws_exception_handler(websocket: WebSocket, exc: CustomWSException) -> None:
     anyio.from_thread.run(websocket.close, status.WS_1013_TRY_AGAIN_LATER)
 
 
@@ -99,10 +100,10 @@ def custom_ws_exception_handler(websocket: WebSocket, exc: CustomWSException):
 )
 
 exception_handlers = {
-    500: error_500,
-    405: method_not_allowed,
-    HTTPException: http_exception,
-    CustomWSException: custom_ws_exception_handler,
+    500: ExceptionHandler(error_500),
+    405: ExceptionHandler(method_not_allowed),
+    HTTPException: ExceptionHandler(http_exception),
+    CustomWSException: ExceptionHandler(custom_ws_exception_handler),
 }
 
 middleware = [
@@ -491,12 +492,14 @@ async def startup():
         assert len(record) == 1
 
 
-def test_middleware_stack_init(test_client_factory: Callable[[ASGIApp], httpx.Client]):
+def test_middleware_stack_init(
+    test_client_factory: Callable[[ASGIApp], httpx.Client]
+) -> None:
     class NoOpMiddleware:
         def __init__(self, app: ASGIApp):
             self.app = app
 
-        async def __call__(self, *args: Any):
+        async def __call__(self, *args: Any) -> None:
             await self.app(*args)
 
     class SimpleInitializableMiddleware:
@@ -506,7 +509,7 @@ def __init__(self, app: ASGIApp):
             self.app = app
             SimpleInitializableMiddleware.counter += 1
 
-        async def __call__(self, *args: Any):
+        async def __call__(self, *args: Any) -> None:
             await self.app(*args)
 
     def get_app() -> ASGIApp: