asgi.py 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179
  1. import typing
  2. import sniffio
  3. from .._models import Request, Response
  4. from .._types import AsyncByteStream
  5. from .base import AsyncBaseTransport
  6. if typing.TYPE_CHECKING: # pragma: no cover
  7. import asyncio
  8. import trio
  9. Event = typing.Union[asyncio.Event, trio.Event]
  10. _Message = typing.Dict[str, typing.Any]
  11. _Receive = typing.Callable[[], typing.Awaitable[_Message]]
  12. _Send = typing.Callable[
  13. [typing.Dict[str, typing.Any]], typing.Coroutine[None, None, None]
  14. ]
  15. _ASGIApp = typing.Callable[
  16. [typing.Dict[str, typing.Any], _Receive, _Send], typing.Coroutine[None, None, None]
  17. ]
  18. def create_event() -> "Event":
  19. if sniffio.current_async_library() == "trio":
  20. import trio
  21. return trio.Event()
  22. else:
  23. import asyncio
  24. return asyncio.Event()
  25. class ASGIResponseStream(AsyncByteStream):
  26. def __init__(self, body: typing.List[bytes]) -> None:
  27. self._body = body
  28. async def __aiter__(self) -> typing.AsyncIterator[bytes]:
  29. yield b"".join(self._body)
  30. class ASGITransport(AsyncBaseTransport):
  31. """
  32. A custom AsyncTransport that handles sending requests directly to an ASGI app.
  33. The simplest way to use this functionality is to use the `app` argument.
  34. ```
  35. client = httpx.AsyncClient(app=app)
  36. ```
  37. Alternatively, you can setup the transport instance explicitly.
  38. This allows you to include any additional configuration arguments specific
  39. to the ASGITransport class:
  40. ```
  41. transport = httpx.ASGITransport(
  42. app=app,
  43. root_path="/submount",
  44. client=("1.2.3.4", 123)
  45. )
  46. client = httpx.AsyncClient(transport=transport)
  47. ```
  48. Arguments:
  49. * `app` - The ASGI application.
  50. * `raise_app_exceptions` - Boolean indicating if exceptions in the application
  51. should be raised. Default to `True`. Can be set to `False` for use cases
  52. such as testing the content of a client 500 response.
  53. * `root_path` - The root path on which the ASGI application should be mounted.
  54. * `client` - A two-tuple indicating the client IP and port of incoming requests.
  55. ```
  56. """
  57. def __init__(
  58. self,
  59. app: _ASGIApp,
  60. raise_app_exceptions: bool = True,
  61. root_path: str = "",
  62. client: typing.Tuple[str, int] = ("127.0.0.1", 123),
  63. ) -> None:
  64. self.app = app
  65. self.raise_app_exceptions = raise_app_exceptions
  66. self.root_path = root_path
  67. self.client = client
  68. async def handle_async_request(
  69. self,
  70. request: Request,
  71. ) -> Response:
  72. assert isinstance(request.stream, AsyncByteStream)
  73. # ASGI scope.
  74. scope = {
  75. "type": "http",
  76. "asgi": {"version": "3.0"},
  77. "http_version": "1.1",
  78. "method": request.method,
  79. "headers": [(k.lower(), v) for (k, v) in request.headers.raw],
  80. "scheme": request.url.scheme,
  81. "path": request.url.path,
  82. "raw_path": request.url.raw_path.split(b"?")[0],
  83. "query_string": request.url.query,
  84. "server": (request.url.host, request.url.port),
  85. "client": self.client,
  86. "root_path": self.root_path,
  87. }
  88. # Request.
  89. request_body_chunks = request.stream.__aiter__()
  90. request_complete = False
  91. # Response.
  92. status_code = None
  93. response_headers = None
  94. body_parts = []
  95. response_started = False
  96. response_complete = create_event()
  97. # ASGI callables.
  98. async def receive() -> typing.Dict[str, typing.Any]:
  99. nonlocal request_complete
  100. if request_complete:
  101. await response_complete.wait()
  102. return {"type": "http.disconnect"}
  103. try:
  104. body = await request_body_chunks.__anext__()
  105. except StopAsyncIteration:
  106. request_complete = True
  107. return {"type": "http.request", "body": b"", "more_body": False}
  108. return {"type": "http.request", "body": body, "more_body": True}
  109. async def send(message: typing.Dict[str, typing.Any]) -> None:
  110. nonlocal status_code, response_headers, response_started
  111. if message["type"] == "http.response.start":
  112. assert not response_started
  113. status_code = message["status"]
  114. response_headers = message.get("headers", [])
  115. response_started = True
  116. elif message["type"] == "http.response.body":
  117. assert not response_complete.is_set()
  118. body = message.get("body", b"")
  119. more_body = message.get("more_body", False)
  120. if body and request.method != "HEAD":
  121. body_parts.append(body)
  122. if not more_body:
  123. response_complete.set()
  124. try:
  125. await self.app(scope, receive, send)
  126. except Exception: # noqa: PIE-786
  127. if self.raise_app_exceptions:
  128. raise
  129. response_complete.set()
  130. if status_code is None:
  131. status_code = 500
  132. if response_headers is None:
  133. response_headers = {}
  134. assert response_complete.is_set()
  135. assert status_code is not None
  136. assert response_headers is not None
  137. stream = ASGIResponseStream(body_parts)
  138. return Response(status_code, headers=response_headers, stream=stream)