Define MiddlewareType to match both classes and callables

Author: Viicos

starlette/applications.py

@@ -9,7 +9,7 @@
 from starlette.requests import Request
 from starlette.responses import Response
 from starlette.routing import BaseRoute, Router
-from starlette.types import ASGIApp, Lifespan, Receive, Scope, Send
+from starlette.types import ASGIApp, Lifespan, MiddlewareType, Receive, Scope, Send
 
 AppType = typing.TypeVar("AppType", bound="Starlette")
 
@@ -135,7 +135,7 @@ def host(
         self.router.host(host, app=app, name=name)
 
     def add_middleware(
-        self, middleware_class: typing.Type[ASGIApp], **options: typing.Any
+        self, middleware_class: MiddlewareType, **options: typing.Any
     ) -> None:  # pragma: no cover
         self.user_middleware.insert(0, Middleware(middleware_class, **options))
 

starlette/middleware/__init__.py

@@ -1,10 +1,10 @@
 import typing
 
-from starlette.types import ASGIApp
+from starlette.types import MiddlewareType
 
 
 class Middleware:
-    def __init__(self, cls: typing.Type[ASGIApp], **options: typing.Any) -> None:
+    def __init__(self, cls: MiddlewareType, **options: typing.Any) -> None:
         self.cls = cls
         self.options = options
 

starlette/types.py

@@ -1,5 +1,11 @@
+import sys
 import typing
 
+if sys.version_info < (3, 8):  # pragma: no cover
+    from typing_extensions import Protocol
+else:  # pragma: no cover
+    from typing import Protocol
+
 AppType = typing.TypeVar("AppType")
 
 Scope = typing.MutableMapping[str, typing.Any]
@@ -15,3 +21,15 @@
     [AppType], typing.AsyncContextManager[typing.Mapping[str, typing.Any]]
 ]
 Lifespan = typing.Union[StatelessLifespan[AppType], StatefulLifespan[AppType]]
+
+
+# This callable protocol can both be used to represent a function returning
+# an ASGIApp, or a class with an __init__ method matching this __call__ signature
+# and a __call__ method matching the ASGIApp signature.
+class MiddlewareType(Protocol):
+    __name__: str
+
+    def __call__(
+        self, *args: typing.Any, **kwargs: typing.Any
+    ) -> ASGIApp:  # pragma: no cover
+        ...

tests/middleware/test_base.py

@@ -1,5 +1,4 @@
 import contextvars
-import typing
 from contextlib import AsyncExitStack
 
 import anyio
@@ -11,7 +10,7 @@
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.responses import PlainTextResponse, StreamingResponse
 from starlette.routing import Route, WebSocketRoute
-from starlette.types import ASGIApp, Receive, Scope, Send
+from starlette.types import ASGIApp, MiddlewareType, Receive, Scope, Send
 
 
 class CustomMiddleware(BaseHTTPMiddleware):
@@ -194,7 +193,7 @@ async def dispatch(self, request, call_next):
         ),
     ],
 )
-def test_contextvars(test_client_factory, middleware_cls: typing.Type[ASGIApp]):
+def test_contextvars(test_client_factory, middleware_cls: MiddlewareType):
     # this has to be an async endpoint because Starlette calls run_in_threadpool
     # on sync endpoints which has it's own set of peculiarities w.r.t propagating
     # contextvars (it propagates them forwards but not backwards)