rewritable_starlette_static.py

"""A module for serving static files with URL rewriting support.

Module Overview:
    This module provides the RewrittenStaticFiles class, an extension of Starlette's StaticFiles that supports URL rewriting.
    The class is designed to serve static files from a specified directory or packages and to apply rewrite rules in case a requested
    file is not found (HTTP 404). Rewrite rules are defined as a dictionary where keys are patterns (either wildcard strings or compiled
    regular expressions) and values are target paths. If a request initially results in a 404 error, the class will attempt to match
    the path against the rewrite rules and serve the corresponding rewritten file if found.

Examples:
    >>> app.mount(
    ...     "/",
    ...     RewrittenStaticFiles(
    ...         directory="./ui/dist",
    ...         html=True,
    ...         rewrite_rules={"favicon.ico": "vite.svg", "*": "index.html"},
    ...     ),
    ...     name="static",
    ... )

"""  # noqa:E501

import fnmatch
import re
from typing import override

from starlette.exceptions import HTTPException
from starlette.responses import Response
from starlette.staticfiles import PathLike, StaticFiles
from starlette.status import HTTP_404_NOT_FOUND
from starlette.types import Scope


class RewrittenStaticFiles(StaticFiles):
    def __init__(  # noqa: PLR0913
        self,
        *,
        directory: PathLike | None = None,
        packages: list[str | tuple[str, str]] | None = None,
        rewrite_rules: dict[str | re.Pattern[str], str] | None = None,
        html: bool = False,
        check_dir: bool = True,
        follow_symlink: bool = False,
    ) -> None:
        super().__init__(
            directory=directory,
            packages=packages,
            html=html,
            check_dir=check_dir,
            follow_symlink=follow_symlink,
        )
        self.rewrite_rules = (
            [
                (
                    p
                    if isinstance(p, re.Pattern)
                    else re.compile(fnmatch.translate(p)),
                    d,
                )
                for p, d in rewrite_rules.items()
            ]
            if rewrite_rules
            else None
        )

    @override
    async def get_response(self, path: str, scope: Scope) -> Response:
        try:
            return await super().get_response(path, scope)
        except HTTPException as exc:
            if exc.status_code != HTTP_404_NOT_FOUND or not (
                self.rewrite_rules
                and (
                    matched_rules := [d for p, d in self.rewrite_rules if p.match(path)]
                )
            ):
                raise
            orig_exc = exc
        for target in matched_rules:
            try:
                return await super().get_response(target, scope)
            except HTTPException as exc:
                if exc.status_code != HTTP_404_NOT_FOUND:
                    raise orig_exc from exc
        raise orig_exc