AsgiWhiteNoise and async WhiteNoiseMiddleware

setup.cfg

@@ -43,6 +43,8 @@ zip_safe = False
 [options.extras_require]
 brotli =
     Brotli
+asgi =
+    aiofile >=3.0, <4.0
 
 [options.packages.find]
 where = src

src/whitenoise/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
-from .base import WhiteNoise
+from .asgi import AsgiWhiteNoise
+from .wsgi import WhiteNoise
 
-__all__ = ["WhiteNoise"]
+__all__ = ["AsgiWhiteNoise", "WhiteNoise"]

src/whitenoise/asgi.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from whitenoise.base import BaseWhiteNoise
+
+
+class AsgiWhiteNoise(BaseWhiteNoise):
+    # This is the same block size as wsgiref.FileWrapper
+    BLOCK_SIZE = 8192
+
+    async def __call__(self, scope, receive, send):
+        static_file = None
+        if scope["type"] == "http":
+            if self.autorefresh:
+                static_file = self.find_file(scope["path"])
+            else:
+                static_file = self.files.get(scope["path"])
+        if static_file is None:
+            await self.application(scope, receive, send)
+        else:
+            await receive(receive)
+            request_headers = convert_asgi_headers(scope["headers"])
+            await self.serve(
+                send, static_file, scope["method"], request_headers, self.BLOCK_SIZE
+            )
+
+    @staticmethod
+    async def serve(send, static_file, method, request_headers, block_size):
+        response = static_file.get_response(method, request_headers)
+        try:
+            await send(
+                {
+                    "type": "http.response.start",
+                    "status": response.status.value,
+                    "headers": convert_wsgi_headers(response.headers),
+                }
+            )
+            if response.file:
+                # We need to only read content-length bytes instead of the whole file,
+                # the difference is important when serving range requests.
+                content_length = int(dict(response.headers)["Content-Length"])
+                for block in read_file(response.file, content_length, block_size):
+                    # TODO: Recode this when ASGI webservers to support zero-copy send
+                    # See https://asgi.readthedocs.io/en/latest/extensions.html#zero-copy-send
+                    await send(
+                        {"type": "http.response.body", "body": block, "more_body": True}
+                    )
+            await send({"type": "http.response.body"})
+        finally:
+            if response.file:
+                response.file.close()
+
+    @staticmethod
+    async def receive(receive):
+        more_body = True
+        while more_body:
+            event = await receive()
+            if event["type"] != "http.request":
+                raise RuntimeError(
+                    "Unexpected ASGI event {!r}, expected {!r}".format(
+                        event["type"], "http.request"
+                    )
+                )
+            more_body = event.get("more_body", False)
+
+
+def read_file(file_handle, content_length, block_size):
+    bytes_left = content_length
+    while bytes_left > 0:
+        data = file_handle.read(min(block_size, bytes_left))
+        if data == b"":
+            raise RuntimeError(
+                f"Premature end of file, expected {bytes_left} more bytes"
+            )
+        bytes_left -= len(data)
+        yield data
+
+
+def convert_asgi_headers(headers):
+    return {
+        "HTTP_" + name.decode().upper().replace("-", "_"): value.decode()
+        for name, value in headers
+    }
+
+
+def convert_wsgi_headers(headers):
+    return [(key.lower().encode(), value.encode()) for key, value in headers]

src/whitenoise/base.py

@@ -5,18 +5,13 @@
 import warnings
 from posixpath import normpath
 from wsgiref.headers import Headers
-from wsgiref.util import FileWrapper
 
 from .media_types import MediaTypes
-from .responders import IsDirectoryError, MissingFileError, Redirect, StaticFile
-from .string_utils import (
-    decode_if_byte_string,
-    decode_path_info,
-    ensure_leading_trailing_slash,
-)
+from .responders import MissingFileError, Redirect, StaticFile
+from .string_utils import decode_if_byte_string, ensure_leading_trailing_slash
 
 
-class WhiteNoise:
+class BaseWhiteNoise:
 
     # Ten years is what nginx sets a max age if you use 'expires max;'
     # so we'll follow its lead
@@ -75,28 +70,6 @@ def __init__(self, application, root=None, prefix=None, **kwargs):
         if root is not None:
             self.add_files(root, prefix)
 
-    def __call__(self, environ, start_response):
-        path = decode_path_info(environ.get("PATH_INFO", ""))
-        if self.autorefresh:
-            static_file = self.find_file(path)
-        else:
-            static_file = self.files.get(path)
-        if static_file is None:
-            return self.application(environ, start_response)
-        else:
-            return self.serve(static_file, environ, start_response)
-
-    @staticmethod
-    def serve(static_file, environ, start_response):
-        response = static_file.get_response(environ["REQUEST_METHOD"], environ)
-        status_line = f"{response.status} {response.status.phrase}"
-        start_response(status_line, list(response.headers))
-        if response.file is not None:
-            file_wrapper = environ.get("wsgi.file_wrapper", FileWrapper)
-            return file_wrapper(response.file)
-        else:
-            return []
-
     def add_files(self, root, prefix=None):
         root = decode_if_byte_string(root, force_text=True)
         root = os.path.abspath(root)

src/whitenoise/middleware.py

@@ -10,8 +10,8 @@
 from django.http import FileResponse
 from django.urls import get_script_prefix
 
-from .base import WhiteNoise
 from .string_utils import decode_if_byte_string, ensure_leading_trailing_slash
+from .wsgi import WhiteNoise
 
 __all__ = ["WhiteNoiseMiddleware"]
 

src/whitenoise/wsgi.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from wsgiref.util import FileWrapper
+
+from .base import BaseWhiteNoise
+from .string_utils import decode_path_info
+
+
+class WhiteNoise(BaseWhiteNoise):
+    def __call__(self, environ, start_response):
+        path = decode_path_info(environ.get("PATH_INFO", ""))
+        if self.autorefresh:
+            static_file = self.find_file(path)
+        else:
+            static_file = self.files.get(path)
+        if static_file is None:
+            return self.application(environ, start_response)
+        else:
+            return self.serve(static_file, environ, start_response)
+
+    @staticmethod
+    def serve(static_file, environ, start_response):
+        response = static_file.get_response(environ["REQUEST_METHOD"], environ)
+        status_line = f"{response.status} {response.status.phrase}"
+        start_response(status_line, list(response.headers))
+        if response.file is not None:
+            file_wrapper = environ.get("wsgi.file_wrapper", FileWrapper)
+            return file_wrapper(response.file)
+        else:
+            return []