ranged_response.py

# Taken from https://gist.github.com/kevinastone/a6a62db57577b3f24e8a6865ed311463
# Context: https://github.com/encode/starlette/pull/1090
from __future__ import annotations

import os
import re
import stat
from typing import NamedTuple
from urllib.parse import quote

import aiofiles
from aiofiles.os import stat as aio_stat
from starlette.datastructures import Headers
from starlette.exceptions import HTTPException
from starlette.responses import Response, guess_type
from starlette.staticfiles import StaticFiles
from starlette.types import Receive, Scope, Send

RANGE_REGEX = re.compile(r"^bytes=(?P<start>\d+)-(?P<end>\d*)$")


class ClosedRange(NamedTuple):
    start: int
    end: int

    def __len__(self) -> int:
        return self.end - self.start + 1

    def __bool__(self) -> bool:
        return len(self) > 0


class OpenRange(NamedTuple):
    start: int
    end: int | None = None

    def clamp(self, start: int, end: int) -> ClosedRange:
        begin = max(self.start, start)
        end = min(x for x in (self.end, end) if x)

        begin = min(begin, end)
        end = max(begin, end)

        return ClosedRange(begin, end)


class RangedFileResponse(Response):
    chunk_size = 4096

    def __init__(
        self,
        path: str | os.PathLike,
        range: OpenRange,
        headers: dict[str, str] | None = None,
        media_type: str | None = None,
        filename: str | None = None,
        stat_result: os.stat_result | None = None,
        method: str | None = None,
    ) -> None:
        if aiofiles is None:
            raise ModuleNotFoundError(
                "'aiofiles' must be installed to use FileResponse"
            )
        self.path = path
        self.range = range
        self.filename = filename
        self.background = None
        self.send_header_only = method is not None and method.upper() == "HEAD"
        if media_type is None:
            media_type = guess_type(filename or path)[0] or "text/plain"
        self.media_type = media_type
        self.init_headers(headers or {})
        if self.filename is not None:
            content_disposition_filename = quote(self.filename)
            if content_disposition_filename != self.filename:
                content_disposition = (
                    f"attachment; filename*=utf-8''{content_disposition_filename}"
                )
            else:
                content_disposition = f'attachment; filename="{self.filename}"'
            self.headers.setdefault("content-disposition", content_disposition)
        self.stat_result = stat_result

    def set_range_headers(self, range: ClosedRange) -> None:
        if not self.stat_result:
            raise ValueError("No stat result to set range headers with")
        total_length = self.stat_result.st_size
        content_length = len(range)
        self.headers["content-range"] = (
            f"bytes {range.start}-{range.end}/{total_length}"
        )
        self.headers["content-length"] = str(content_length)
        pass

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:  # noqa: ARG002
        if self.stat_result is None:
            try:
                stat_result = await aio_stat(self.path)
                self.stat_result = stat_result
            except FileNotFoundError as fnfe:
                raise RuntimeError(
                    f"File at path {self.path} does not exist."
                ) from fnfe
            else:
                mode = stat_result.st_mode
                if not stat.S_ISREG(mode):
                    raise RuntimeError(f"File at path {self.path} is not a file.")

        byte_range = self.range.clamp(0, self.stat_result.st_size)
        self.set_range_headers(byte_range)

        async with aiofiles.open(self.path, mode="rb") as file:
            await file.seek(byte_range.start)
            await send(
                {
                    "type": "http.response.start",
                    "status": 206,
                    "headers": self.raw_headers,
                }
            )
            if self.send_header_only:
                await send(
                    {"type": "http.response.body", "body": b"", "more_body": False}
                )
            else:
                remaining_bytes = len(byte_range)

                if not byte_range:
                    await send(
                        {"type": "http.response.body", "body": b"", "more_body": False}
                    )
                    return

                while remaining_bytes > 0:
                    chunk_size = min(self.chunk_size, remaining_bytes)
                    chunk = await file.read(chunk_size)
                    remaining_bytes -= len(chunk)
                    await send(
                        {
                            "type": "http.response.body",
                            "body": chunk,
                            "more_body": remaining_bytes > 0,
                        }
                    )


class RangedStaticFiles(StaticFiles):
    def file_response(
        self,
        full_path: str | os.PathLike,
        stat_result: os.stat_result,
        scope: Scope,
        status_code: int = 200,
    ) -> Response:
        request_headers = Headers(scope=scope)

        if request_headers.get("range"):
            response = self.ranged_file_response(
                full_path, stat_result=stat_result, scope=scope
            )
        else:
            response = super().file_response(
                full_path, stat_result=stat_result, scope=scope, status_code=status_code
            )
        response.headers["accept-ranges"] = "bytes"
        return response

    def ranged_file_response(
        self,
        full_path: str | os.PathLike,
        stat_result: os.stat_result,
        scope: Scope,
    ) -> Response:
        method = scope["method"]
        request_headers = Headers(scope=scope)

        range_header = request_headers["range"]

        match = RANGE_REGEX.search(range_header)
        if not match:
            raise HTTPException(400)

        start, end = match.group("start"), match.group("end")

        range = OpenRange(int(start), int(end) if end else None)

        return RangedFileResponse(
            full_path, range, stat_result=stat_result, method=method
        )