Mabasej_Team/.venv/lib/python3.9/site-packages/starlette/middleware/wsgi.py

144 lines
5.1 KiB
Python
Raw Normal View History

2021-03-17 08:57:57 +01:00
import asyncio
import io
import sys
import typing
from starlette.concurrency import run_in_threadpool
from starlette.types import Message, Receive, Scope, Send
def build_environ(scope: Scope, body: bytes) -> dict:
"""
Builds a scope and request body into a WSGI environ object.
"""
environ = {
"REQUEST_METHOD": scope["method"],
"SCRIPT_NAME": scope.get("root_path", "").encode("utf8").decode("latin1"),
"PATH_INFO": scope["path"].encode("utf8").decode("latin1"),
"QUERY_STRING": scope["query_string"].decode("ascii"),
"SERVER_PROTOCOL": f"HTTP/{scope['http_version']}",
"wsgi.version": (1, 0),
"wsgi.url_scheme": scope.get("scheme", "http"),
"wsgi.input": io.BytesIO(body),
"wsgi.errors": sys.stdout,
"wsgi.multithread": True,
"wsgi.multiprocess": True,
"wsgi.run_once": False,
}
# Get server name and port - required in WSGI, not in ASGI
server = scope.get("server") or ("localhost", 80)
environ["SERVER_NAME"] = server[0]
environ["SERVER_PORT"] = server[1]
# Get client IP address
if scope.get("client"):
environ["REMOTE_ADDR"] = scope["client"][0]
# Go through headers and make them into environ entries
for name, value in scope.get("headers", []):
name = name.decode("latin1")
if name == "content-length":
corrected_name = "CONTENT_LENGTH"
elif name == "content-type":
corrected_name = "CONTENT_TYPE"
else:
corrected_name = f"HTTP_{name}".upper().replace("-", "_")
# HTTPbis say only ASCII chars are allowed in headers, but we latin1 just in case
value = value.decode("latin1")
if corrected_name in environ:
value = environ[corrected_name] + "," + value
environ[corrected_name] = value
return environ
class WSGIMiddleware:
def __init__(self, app: typing.Callable, workers: int = 10) -> None:
self.app = app
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
assert scope["type"] == "http"
responder = WSGIResponder(self.app, scope)
await responder(receive, send)
class WSGIResponder:
def __init__(self, app: typing.Callable, scope: Scope) -> None:
self.app = app
self.scope = scope
self.status = None
self.response_headers = None
self.send_event = asyncio.Event()
self.send_queue = [] # type: typing.List[typing.Optional[Message]]
self.loop = asyncio.get_event_loop()
self.response_started = False
self.exc_info = None # type: typing.Any
async def __call__(self, receive: Receive, send: Send) -> None:
body = b""
more_body = True
while more_body:
message = await receive()
body += message.get("body", b"")
more_body = message.get("more_body", False)
environ = build_environ(self.scope, body)
sender = None
try:
sender = self.loop.create_task(self.sender(send))
await run_in_threadpool(self.wsgi, environ, self.start_response)
self.send_queue.append(None)
self.send_event.set()
await asyncio.wait_for(sender, None)
if self.exc_info is not None:
raise self.exc_info[0].with_traceback(
self.exc_info[1], self.exc_info[2]
)
finally:
if sender and not sender.done():
sender.cancel() # pragma: no cover
async def sender(self, send: Send) -> None:
while True:
if self.send_queue:
message = self.send_queue.pop(0)
if message is None:
return
await send(message)
else:
await self.send_event.wait()
self.send_event.clear()
def start_response(
self,
status: str,
response_headers: typing.List[typing.Tuple[str, str]],
exc_info: typing.Any = None,
) -> None:
self.exc_info = exc_info
if not self.response_started:
self.response_started = True
status_code_string, _ = status.split(" ", 1)
status_code = int(status_code_string)
headers = [
(name.strip().encode("ascii").lower(), value.strip().encode("ascii"))
for name, value in response_headers
]
self.send_queue.append(
{
"type": "http.response.start",
"status": status_code,
"headers": headers,
}
)
self.loop.call_soon_threadsafe(self.send_event.set)
def wsgi(self, environ: dict, start_response: typing.Callable) -> None:
for chunk in self.app(environ, start_response):
self.send_queue.append(
{"type": "http.response.body", "body": chunk, "more_body": True}
)
self.loop.call_soon_threadsafe(self.send_event.set)
self.send_queue.append({"type": "http.response.body", "body": b""})
self.loop.call_soon_threadsafe(self.send_event.set)