Staticfiles handle 404 (#54)

.github/workflows/ci.yml

@@ -47,7 +47,7 @@ jobs:
         run: pdm run check
 
       - name: Tests
-        run: pdm run test --cov=./baize -o log_cli=true -o log_cli_level=DEBUG
+        run: pdm run test --cov=./baize --cov-report=xml -o log_cli=true -o log_cli_level=DEBUG
 
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v3

baize/asgi/staticfiles.py

@@ -1,23 +1,38 @@
+import os
 import stat
 
 from baize import staticfiles
 from baize.datastructures import URL
 from baize.exceptions import HTTPException
-from baize.typing import Receive, Scope, Send
+from baize.typing import ASGIApp, Receive, Scope, Send
 
 from .responses import FileResponse, RedirectResponse, Response
 
 
-class Files(staticfiles.BaseFiles):
+class Files(staticfiles.BaseFiles[ASGIApp]):
     """
     Provide the ASGI application to download files in the specified path or
     the specified directory under the specified package.
 
     Support request range and cache (304 status code).
-
-    NOTE: Need users handle HTTPException(404).
     """
 
+    def file_response(
+        self,
+        filepath: str,
+        stat_result: os.stat_result,
+        if_none_match: str,
+        if_modified_since: str,
+    ) -> Response:
+        if self.if_none_match(
+            FileResponse.generate_etag(stat_result), if_none_match
+        ) or self.if_modified_since(stat_result.st_ctime, if_modified_since):
+            response = Response(304)
+        else:
+            response = FileResponse(filepath, stat_result=stat_result)
+            self.set_response_headers(response)
+        return response
+
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if_none_match: str = ""
         if_modified_since: str = ""
@@ -30,28 +45,24 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         stat_result, is_file = self.check_path_is_file(filepath)
         if is_file and stat_result:
             assert filepath is not None  # Just for type check
-            if self.if_none_match(
-                FileResponse.generate_etag(stat_result), if_none_match
-            ) or self.if_modified_since(stat_result.st_ctime, if_modified_since):
-                response = Response(304)
-            else:
-                response = FileResponse(filepath, stat_result=stat_result)
-            self.set_response_headers(response)
-            return await response(scope, receive, send)
+            return await self.file_response(
+                filepath, stat_result, if_none_match, if_modified_since
+            )(scope, receive, send)
 
-        raise HTTPException(404)
+        if self.handle_404 is None:
+            raise HTTPException(404)
+        else:
+            return await self.handle_404(scope, receive, send)
 
 
-class Pages(staticfiles.BasePages):
+class Pages(staticfiles.BasePages[ASGIApp], Files):
     """
     Provide the ASGI application to download files in the specified path or
     the specified directory under the specified package.
     Unlike `Files`, when you visit a directory, it will try to return the content
     of the file named `index.html` in that directory.
 
     Support request range and cache (304 status code).
-
-    NOTE: Need users handle HTTPException(404).
     """
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
@@ -75,17 +86,15 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if stat_result is not None:
             assert filepath is not None  # Just for type check
             if is_file:
-                if self.if_none_match(
-                    FileResponse.generate_etag(stat_result), if_none_match
-                ) or self.if_modified_since(stat_result.st_ctime, if_modified_since):
-                    response = Response(304)
-                else:
-                    response = FileResponse(filepath, stat_result=stat_result)
-                self.set_response_headers(response)
-                return await response(scope, receive, send)
+                return await self.file_response(
+                    filepath, stat_result, if_none_match, if_modified_since
+                )(scope, receive, send)
             if stat.S_ISDIR(stat_result.st_mode):
                 url = URL(scope=scope)
                 url = url.replace(scheme="", path=url.path + "/")
                 return await RedirectResponse(url)(scope, receive, send)
 
-        raise HTTPException(404)
+        if self.handle_404 is None:
+            raise HTTPException(404)
+        else:
+            return await self.handle_404(scope, receive, send)

baize/staticfiles.py

@@ -2,10 +2,10 @@
 import os
 import stat
 from email.utils import parsedate_to_datetime
-from typing import Optional, Tuple, Union
+from typing import Generic, Optional, Tuple, TypeVar, Union
 
 from .responses import BaseResponse
-from .typing import Literal
+from .typing import ASGIApp, Literal, WSGIApp
 
 try:
     from mypy_extensions import mypyc_attr
@@ -15,20 +15,25 @@ def mypyc_attr(*attrs, **kwattrs):  # type: ignore
         return lambda x: x
 
 
+Interface = TypeVar("Interface", ASGIApp, WSGIApp)
+
+
 @mypyc_attr(allow_interpreted_subclasses=True)
-class BaseFiles:
+class BaseFiles(Generic[Interface]):
     def __init__(
         self,
         directory: Union[str, "os.PathLike[str]"],
         package: Optional[str] = None,
         *,
+        handle_404: Optional[Interface] = None,
         cacheability: Literal["public", "private", "no-cache", "no-store"] = "public",
         max_age: int = 60 * 10,  # 10 minutes
     ) -> None:
         assert not (
             os.path.isabs(directory) and package is not None
         ), "directory must be a relative path, with package is not None"
         self.directory = self.normalize_dir_path(str(directory), package)
+        self.handle_404: Optional[Interface] = handle_404
         self.cacheability = cacheability
         self.max_age = max_age
 
@@ -103,7 +108,7 @@ def set_response_headers(self, response: BaseResponse) -> None:
 
 
 @mypyc_attr(allow_interpreted_subclasses=True)
-class BasePages(BaseFiles):
+class BasePages(BaseFiles[Interface], Generic[Interface]):
     def ensure_absolute_path(self, path: str) -> Optional[str]:
         abspath = super().ensure_absolute_path(path)
         if abspath is not None:

baize/wsgi/staticfiles.py

@@ -1,24 +1,39 @@
+import os
 import stat
 from typing import Iterable
 
 from baize import staticfiles
 from baize.datastructures import URL
 from baize.exceptions import HTTPException
-from baize.typing import Environ, StartResponse
+from baize.typing import Environ, StartResponse, WSGIApp
 
 from .responses import FileResponse, RedirectResponse, Response
 
 
-class Files(staticfiles.BaseFiles):
+class Files(staticfiles.BaseFiles[WSGIApp]):
     """
     Provide the WSGI application to download files in the specified path or
     the specified directory under the specified package.
 
     Support request range and cache (304 status code).
-
-    NOTE: Need users handle HTTPException(404).
     """
 
+    def file_response(
+        self,
+        filepath: str,
+        stat_result: os.stat_result,
+        if_none_match: str,
+        if_modified_since: str,
+    ) -> Response:
+        if self.if_none_match(
+            FileResponse.generate_etag(stat_result), if_none_match
+        ) or self.if_modified_since(stat_result.st_ctime, if_modified_since):
+            response = Response(304)
+        else:
+            response = FileResponse(filepath, stat_result=stat_result)
+        self.set_response_headers(response)
+        return response
+
     def __call__(
         self, environ: Environ, start_response: StartResponse
     ) -> Iterable[bytes]:
@@ -28,19 +43,17 @@ def __call__(
         stat_result, is_file = self.check_path_is_file(filepath)
         if is_file and stat_result:
             assert filepath is not None  # Just for type check
-            if self.if_none_match(
-                FileResponse.generate_etag(stat_result), if_none_match
-            ) or self.if_modified_since(stat_result.st_ctime, if_modified_since):
-                response = Response(304)
-            else:
-                response = FileResponse(filepath, stat_result=stat_result)
-            self.set_response_headers(response)
-            return response(environ, start_response)
+            return self.file_response(
+                filepath, stat_result, if_none_match, if_modified_since
+            )(environ, start_response)
 
-        raise HTTPException(404)
+        if self.handle_404 is None:
+            raise HTTPException(404)
+        else:
+            return self.handle_404(environ, start_response)
 
 
-class Pages(staticfiles.BasePages):
+class Pages(staticfiles.BasePages[WSGIApp], Files):
     """
     Provide the WSGI application to download files in the specified path or
     the specified directory under the specified package.
@@ -49,8 +62,6 @@ class Pages(staticfiles.BasePages):
     exist, it will return the content of that file.
 
     Support request range and cache (304 status code).
-
-    NOTE: Need users handle HTTPException(404).
     """
 
     def __call__(
@@ -71,17 +82,15 @@ def __call__(
         if stat_result is not None:
             assert filepath is not None  # Just for type check
             if is_file:
-                if self.if_none_match(
-                    FileResponse.generate_etag(stat_result), if_none_match
-                ) or self.if_modified_since(stat_result.st_ctime, if_modified_since):
-                    response = Response(304)
-                else:
-                    response = FileResponse(filepath, stat_result=stat_result)
-                self.set_response_headers(response)
-                return response(environ, start_response)
+                return self.file_response(
+                    filepath, stat_result, if_none_match, if_modified_since
+                )(environ, start_response)
             if stat.S_ISDIR(stat_result.st_mode):
                 url = URL(environ=environ)
                 url = url.replace(scheme="", path=url.path + "/")
                 return RedirectResponse(url)(environ, start_response)
 
-        raise HTTPException(404)
+        if self.handle_404 is None:
+            raise HTTPException(404)
+        else:
+            return self.handle_404(environ, start_response)

tests/test_asgi.py

@@ -1067,6 +1067,7 @@ async def test_hosts():
     [
         Files(Path(__file__).absolute().parent.parent / "baize"),
         Files(".", "baize"),
+        Files(".", "baize", handle_404=PlainTextResponse("", 404)),
     ],
 )
 async def test_files(app):
@@ -1107,11 +1108,16 @@ async def test_files(app):
             )
         ).status_code == 304
 
-        with pytest.raises(HTTPException):
-            await client.get("/")
+        if app.handle_404 is None:
 
-        with pytest.raises(HTTPException):
-            await client.get("/%2E%2E/baize/%2E%2E/%2E%2E/README.md")
+            with pytest.raises(HTTPException):
+                await client.get("/")
+
+            with pytest.raises(HTTPException):
+                await client.get("/%2E%2E/baize/%2E%2E/%2E%2E/README.md")
+
+        else:
+            assert (await client.get("/")).status_code == 404
 
 
 @pytest.mark.asyncio

tests/test_wsgi.py

@@ -565,6 +565,7 @@ def test_hosts():
     [
         Files(Path(__file__).absolute().parent.parent / "baize"),
         Files(".", "baize"),
+        Files(".", "baize", handle_404=PlainTextResponse("", 404)),
     ],
 )
 def test_files(app):
@@ -603,11 +604,16 @@ def test_files(app):
             )
         ).status_code == 304
 
-        with pytest.raises(HTTPException):
-            client.get("/")
+        if app.handle_404 is None:
 
-        with pytest.raises(HTTPException):
-            client.get("/%2E%2E/baize/%2E%2E/%2E%2E/README.md")
+            with pytest.raises(HTTPException):
+                client.get("/")
+
+            with pytest.raises(HTTPException):
+                client.get("/%2E%2E/baize/%2E%2E/%2E%2E/README.md")
+
+        else:
+            assert client.get("/").status_code == 404
 
 
 def test_pages(tmpdir):