404 page

src/ziplineio/app.py

@@ -1,4 +1,5 @@
 import inspect
+import re
 from typing import Any, Callable, List, Tuple, Type
 
 
@@ -8,7 +9,7 @@
 from ziplineio import settings
 from ziplineio.handler import Handler
 from ziplineio.request import Request
-from ziplineio.response import format_response
+from ziplineio.response import NotFoundResponse, format_response
 from ziplineio.router import Router
 from ziplineio.static import staticfiles
 from ziplineio.utils import call_handler, parse_scope
@@ -37,6 +38,9 @@ def put(self, path: str) -> Callable[[Handler], Callable]:
     def delete(self, path: str) -> Callable[[Handler], Callable]:
         return self._router.delete(path)
 
+    def not_found(self, handler: Handler) -> None:
+        self._router.not_found(handler)
+
     def get_handler(self, method: str, path: str) -> Tuple[Handler, dict]:
         app_level_deps = self._injector.get_injected_services("app")
         handler, params = self._router.get_handler(method, path)
@@ -94,7 +98,15 @@ async def _get_and_call_handler(
         )
 
         # If middleware does not provide a response, return a 404 Not Found
-        return res if res is not None else NotFoundHttpException()
+        # return res if res is not None else NotFoundHttpException()
+        if res is not None:
+            return res
+
+        if self._router._not_found_handler:
+            body = await call_handler(self._router._not_found_handler, req)
+            return NotFoundResponse(body)
+
+        return NotFoundHttpException()
 
     def __call__(self, *args: Any, **kwds: Any) -> Any:
         async def uvicorn_handler(scope: dict, receive: Any, send: Any) -> None:

src/ziplineio/response.py

@@ -14,13 +14,27 @@ class RawResponse(TypedDict):
 class Response:
     def __init__(self, status: int, headers: Dict[str, str], body: str):
         self.status = status
-        self.headers = headers
+        self._headers = headers
         self.body = body
 
     status: int
-    headers: Dict[str, str]
+    _headers: Dict[str, str]
     body: str
 
+    # If content-type is not provided, infer content-type from the body
+    def get_headers(self) -> Dict[str, str]:
+        headers = self._headers.copy()
+        if "Content-Type" not in headers:
+            if isinstance(self.body, bytes):
+                headers["Content-Type"] = "text/plain"
+            elif isinstance(self.body, str):
+                headers["Content-Type"] = "text/html"
+            elif isinstance(self.body, dict):
+                headers["Content-Type"] = "application/json"
+            else:
+                headers["Content-Type"] = "text/plain"
+        return headers
+
     def __len__(self) -> int:
         return 1
 
@@ -40,6 +54,11 @@ def __init__(self, body: str):
         super().__init__(200, {"Content-Type": "text/html"}, body)
 
 
+class NotFoundResponse(Response):
+    def __init__(self, body: str):
+        super().__init__(404, {}, body)
+
+
 def format_headers(headers: Dict[str, str] | None) -> List[Tuple[bytes, bytes]]:
     if headers is None:
         return []
@@ -53,6 +72,8 @@ def format_body(body: bytes | str | dict) -> bytes:
         return bytes(body, "utf-8")
     elif isinstance(body, dict):
         return bytes(json.dumps(body), "utf-8")
+    elif isinstance(body, Response):
+        return format_body(body.body)
     raise ValueError("Invalid body type")
 
 
@@ -85,7 +106,7 @@ def format_response(
         }
     elif isinstance(response, Response):
         raw_response = {
-            "headers": format_headers(response.headers),
+            "headers": format_headers(response.get_headers()),
             "status": response.status,
             "body": format_body(response.body),
         }

src/ziplineio/router.py

@@ -11,6 +11,7 @@ class Router:
     _id: str
     _handlers: Dict[str, Dict[str, Callable]]
     _router_level_middelwares: List[Handler]
+    _not_found_handler: Handler
     _sub_routers: Dict[str, "Router"]
     _prefix: str
     _injector: DependencyInjector
@@ -19,6 +20,7 @@ def __init__(self, prefix: str = "") -> None:
         self._id = str(uuid.uuid4())
         self._handlers = {"GET": {}, "POST": {}, "PUT": {}, "DELETE": {}}
         self._router_level_middelwares = []
+        self._not_found_handler = None
         self._sub_routers = {}
         self._prefix = prefix.rstrip("/")
         self._injector = injector
@@ -64,6 +66,9 @@ def put(self, path: str) -> Callable[[Callable], Callable]:
     def delete(self, path: str) -> Callable[[Callable], Callable]:
         return self.route("DELETE", path)
 
+    def not_found(self, handler: Handler) -> None:
+        self._not_found_handler = handler
+
     def add_sub_router(self, prefix: str, sub_router: "Router") -> None:
         self._sub_routers[prefix.rstrip("/")] = sub_router
 

test/mocks/templates/404.html

@@ -0,0 +1,14 @@
+<section class="bg-indigo-600 text-white py-20">
+  <div class="container mx-auto px-6 text-center">
+    <h1 class="text-7xl font-bold">404</h1>
+    <h2 class="mt-4 text-4xl font-semibold">Page Not Found</h2>
+    <p class="mt-4 text-lg">
+      Sorry, the page you're looking for doesn't exist or has been moved.
+    </p>
+    <a
+      href="/"
+      class="mt-8 inline-block px-6 py-3 bg-white text-indigo-600 font-semibold rounded-lg shadow hover:bg-gray-100"
+      >Go Back Home</a
+    >
+  </div>
+</section>

test/test_e2e.py

@@ -66,6 +66,12 @@ async def user_handler(req):
     return {"message": f"User {req.path_params['id']} received"}
 
 
+@app.not_found
+@jinja(env, "404.html")
+def not_found():
+    return {"current_route": "404"}
+
+
 def run_server():
     uvicorn.run(app, port=5050)
 
@@ -82,9 +88,6 @@ async def test_handler_returns_bytes(self):
         self.assertEqual(response.status_code, 200)
         self.assertEqual(response.content, b"Hello, world!")
 
-        print("BYTES")
-        print(response.content)
-
     async def test_handler_returns_str(self):
         response = requests.get("http://localhost:5050/str")
         self.assertEqual(response.status_code, 200)
@@ -99,14 +102,15 @@ async def test_sync_route(self):
         response = requests.get("http://localhost:5050/sync-thread")
         self.assertEqual(response.status_code, 200)
 
-    async def test_404(self):
-        response = requests.get("http://localhost:5050/some-random-route")
-        self.assertEqual(response.status_code, 404)
-
     async def test_jinja(self):
         response = requests.get("http://localhost:5050/jinja")
         self.assertEqual(response.status_code, 200)
         self.assertTrue("service1 content" in response.text)
 
+    async def test_404_jinja(self):
+        response = requests.get("http://localhost:5050/some-random-route")
+        self.assertEqual(response.status_code, 404)
+        self.assertTrue("404" in response.text)
+
     async def asyncTearDown(self):
         self.proc.terminate()

test/test_html.py

@@ -34,9 +34,6 @@ async def asyncSetUp(self):
     async def test_render_jinja(self):
         response = requests.get("http://localhost:5050/")
 
-        print("Resooinse")
-        print(response.text)
-
         self.assertEqual(response.status_code, 200)
         self.assertEqual(response.headers["Content-Type"], "text/html")
         self.assertTrue("<p>Welcome to the home page!</p>" in response.text)

test/test_middleware.py

@@ -86,8 +86,6 @@ async def test_handler_with_middleware(req: Request, ctx: dict):
         response = await call_handler(handler, req)
         response = format_response(response, settings.DEFAULT_HEADERS)
 
-        print(f"response: {response}")
-
         # Assertions
         self.assertEqual(response["status"], 500)
         self.assertEqual(
@@ -119,8 +117,6 @@ async def test_handler_with_middleware(req: Request, ctx: dict):
         handler, params = self.app._router.get_handler("GET", "/with-middleware")
         response = await call_handler(handler, req)
 
-        print(response)
-
         # Assertions
         self.assertEqual(response["message"], "Hi from middleware 2")
 

test/test_response.py

@@ -35,7 +35,6 @@ async def test_static_file(self):
         self.assertEqual(r.status_code, 200)
         self.assertEqual(r.headers["Content-Type"], "text/css")
         self.assertTrue("background-color: #f0f0f0;" in r.text)
-        print(r.content)
 
     async def asyncTearDown(self):
         self.proc.terminate()