|
1 | 1 | import base64 |
2 | 2 | import functools |
| 3 | +import inspect |
3 | 4 | from abc import ABC, abstractmethod |
4 | 5 | from asyncio import CancelledError, sleep |
5 | | -from collections.abc import AsyncIterator, Iterable, Mapping, Sequence |
| 6 | +from collections.abc import ( |
| 7 | + AsyncGenerator, |
| 8 | + AsyncIterator, |
| 9 | + Callable, |
| 10 | + Iterable, |
| 11 | + Mapping, |
| 12 | + Sequence, |
| 13 | +) |
6 | 14 | from dataclasses import replace |
7 | 15 | from http import HTTPStatus |
8 | | -from typing import TYPE_CHECKING, TypeVar |
| 16 | +from typing import TYPE_CHECKING, Generic, TypeVar, cast |
9 | 17 | from urllib.parse import parse_qs |
10 | 18 |
|
11 | 19 | from . import _compression, _server_shared |
|
48 | 56 | Scope = "asgiref.typing.Scope" |
49 | 57 |
|
50 | 58 |
|
| 59 | +_SVC = TypeVar("_SVC") |
51 | 60 | _REQ = TypeVar("_REQ") |
52 | 61 | _RES = TypeVar("_RES") |
53 | 62 |
|
|
64 | 73 | ) |
65 | 74 |
|
66 | 75 |
|
67 | | -class ConnectASGIApplication(ABC): |
| 76 | +class ConnectASGIApplication(ABC, Generic[_SVC]): |
68 | 77 | """An ASGI application for the Connect protocol.""" |
69 | 78 |
|
| 79 | + _resolved_endpoints: Mapping[str, Endpoint] | None |
| 80 | + |
70 | 81 | @property |
71 | 82 | @abstractmethod |
72 | 83 | def path(self) -> str: ... |
73 | 84 |
|
74 | 85 | def __init__( |
75 | 86 | self, |
76 | 87 | *, |
77 | | - endpoints: Mapping[str, Endpoint], |
| 88 | + service: _SVC | AsyncGenerator[_SVC], |
| 89 | + endpoints: Callable[[_SVC], Mapping[str, Endpoint]], |
78 | 90 | interceptors: Iterable[Interceptor] = (), |
79 | 91 | read_max_bytes: int | None = None, |
80 | 92 | ) -> None: |
81 | 93 | """Initialize the ASGI application.""" |
82 | 94 | super().__init__() |
83 | | - if interceptors: |
84 | | - interceptors = resolve_interceptors(interceptors) |
85 | | - endpoints = { |
86 | | - path: _apply_interceptors(endpoint, interceptors) |
87 | | - for path, endpoint in endpoints.items() |
88 | | - } |
| 95 | + self._service = service |
89 | 96 | self._endpoints = endpoints |
| 97 | + self._interceptors = interceptors |
| 98 | + self._resolved_endpoints = None |
90 | 99 | self._read_max_bytes = read_max_bytes |
91 | 100 |
|
92 | 101 | async def __call__( |
93 | 102 | self, scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable |
94 | 103 | ) -> None: |
95 | | - assert scope["type"] == "http" # noqa: S101 - only for type narrowing, in practice always true |
| 104 | + if scope["type"] == "websocket": |
| 105 | + msg = "connect does not support websockets" |
| 106 | + raise RuntimeError(msg) |
| 107 | + |
| 108 | + if scope["type"] == "lifespan": |
| 109 | + service_iter = None |
| 110 | + while True: |
| 111 | + msg = await receive() |
| 112 | + match msg["type"]: |
| 113 | + case "lifespan.startup": |
| 114 | + # Need to cast since type checking doesn't seem to narrow well with isasyncgen |
| 115 | + if inspect.isasyncgen(self._service): |
| 116 | + service_iter = cast( |
| 117 | + "AsyncGenerator[_SVC, None]", self._service |
| 118 | + ) |
| 119 | + try: |
| 120 | + service = await anext(service_iter) |
| 121 | + except Exception as e: |
| 122 | + await send( |
| 123 | + { |
| 124 | + "type": "lifespan.startup.failed", |
| 125 | + "message": str(e), |
| 126 | + } |
| 127 | + ) |
| 128 | + return None |
| 129 | + else: |
| 130 | + service = cast("_SVC", self._service) |
| 131 | + self._resolved_endpoints = self._resolve_endpoints(service) |
| 132 | + await send({"type": "lifespan.startup.complete"}) |
| 133 | + case "lifespan.shutdown": |
| 134 | + if service_iter is not None: |
| 135 | + try: |
| 136 | + await service_iter.aclose() |
| 137 | + except Exception as e: |
| 138 | + await send( |
| 139 | + { |
| 140 | + "type": "lifespan.shutdown.failed", |
| 141 | + "message": str(e), |
| 142 | + } |
| 143 | + ) |
| 144 | + return None |
| 145 | + await send({"type": "lifespan.shutdown.complete"}) |
| 146 | + return None |
| 147 | + |
| 148 | + if not self._resolved_endpoints: |
| 149 | + if inspect.isasyncgen(self._service): |
| 150 | + msg = "ASGI server does not support lifespan but async generator passed for service. Enable lifespan support." |
| 151 | + raise RuntimeError(msg) |
| 152 | + |
| 153 | + self._resolved_endpoints = self._resolve_endpoints( |
| 154 | + cast("_SVC", self._service) |
| 155 | + ) |
| 156 | + endpoints = self._resolved_endpoints |
96 | 157 |
|
97 | 158 | ctx: RequestContext | None = None |
98 | 159 | try: |
99 | 160 | path = scope["path"] |
100 | | - endpoint = self._endpoints.get(path) |
| 161 | + endpoint = endpoints.get(path) |
101 | 162 | if not endpoint and scope["root_path"]: |
102 | 163 | # The application was mounted at some root so try stripping the prefix. |
103 | 164 | path = path.removeprefix(scope["root_path"]) |
104 | | - endpoint = self._endpoints.get(path) |
105 | | - |
| 165 | + endpoint = endpoints.get(path) |
106 | 166 | if not endpoint: |
107 | 167 | raise HTTPException(HTTPStatus.NOT_FOUND, []) |
108 | 168 |
|
@@ -381,6 +441,17 @@ async def _handle_error( |
381 | 441 | ) |
382 | 442 | await send({"type": "http.response.body", "body": body, "more_body": False}) |
383 | 443 |
|
| 444 | + def _resolve_endpoints(self, service: _SVC) -> Mapping[str, Endpoint]: |
| 445 | + resolved_endpoints = self._endpoints(service) |
| 446 | + if self._interceptors: |
| 447 | + resolved_endpoints = { |
| 448 | + path: _apply_interceptors( |
| 449 | + endpoint, resolve_interceptors(self._interceptors) |
| 450 | + ) |
| 451 | + for path, endpoint in resolved_endpoints.items() |
| 452 | + } |
| 453 | + return resolved_endpoints |
| 454 | + |
384 | 455 |
|
385 | 456 | async def _send_stream_response_headers( |
386 | 457 | send: ASGISendCallable, codec: Codec, compression_name: str, ctx: RequestContext |
|
0 commit comments