diff --git a/robyn/__init__.py b/robyn/__init__.py index 08dd0631d..aad3ce25d 100644 --- a/robyn/__init__.py +++ b/robyn/__init__.py @@ -20,7 +20,7 @@ from robyn.openapi import OpenAPI from robyn.processpool import run_processes from robyn.reloader import compile_rust_files -from robyn.responses import SSEMessage, SSEResponse, StreamingResponse, html, serve_file, serve_html +from robyn.responses import RedirectResponse, SSEMessage, SSEResponse, StreamingResponse, html, serve_file, serve_html from robyn.robyn import FunctionInfo, Headers, HttpMethod, Request, Response, WebSocketConnector, get_version from robyn.router import MiddlewareRouter, MiddlewareType, Router, WebSocketRouter from robyn.types import Directory, JsonBody @@ -802,6 +802,7 @@ def cors_middleware(request): "StreamingResponse", "SSEResponse", "SSEMessage", + "RedirectResponse", "ALLOW_CORS", "SubRouter", "AuthenticationHandler", diff --git a/robyn/responses.py b/robyn/responses.py index 8d3717f11..a74a8433f 100644 --- a/robyn/responses.py +++ b/robyn/responses.py @@ -5,6 +5,8 @@ from robyn.robyn import Headers, Response +_REDIRECT_STATUS_CODES = frozenset({301, 302, 303, 307, 308}) + class FileResponse: def __init__( @@ -19,6 +21,33 @@ def __init__( self.headers = headers or Headers({"Content-Disposition": "attachment"}) +class RedirectResponse(Response): + """Convenience response that issues an HTTP redirect. + + Args: + url: The target URL to redirect to. + status_code: HTTP status code (default 307 Temporary Redirect). + Common values: 301 (permanent), 302 (found), 303 (see other), 307 (temporary), 308 (permanent redirect). + headers: Optional additional headers. + """ + + def __init__( + self, + url: str, + status_code: int = 307, + headers: Optional[Headers] = None, + ): + if status_code not in _REDIRECT_STATUS_CODES: + raise ValueError(f"Invalid redirect status code {status_code}. Must be one of: {sorted(_REDIRECT_STATUS_CODES)}") + redirect_headers = headers or Headers({}) + redirect_headers.set("Location", url) + super().__init__( + status_code=status_code, + headers=redirect_headers, + description="", + ) + + def html(html: str) -> Response: """ This function will help in serving a simple html string diff --git a/unit_tests/test_redirect_response.py b/unit_tests/test_redirect_response.py new file mode 100644 index 000000000..5d316902d --- /dev/null +++ b/unit_tests/test_redirect_response.py @@ -0,0 +1,29 @@ +import pytest + +from robyn.responses import RedirectResponse + + +def test_redirect_response_defaults(): + resp = RedirectResponse("/new-location") + assert resp.status_code == 307 + assert resp.headers.get("Location") == "/new-location" + + +def test_redirect_response_301(): + resp = RedirectResponse("/permanent", status_code=301) + assert resp.status_code == 301 + assert resp.headers.get("Location") == "/permanent" + + +def test_redirect_response_with_extra_headers(): + from robyn.robyn import Headers + + h = Headers({"X-Custom": "value"}) + resp = RedirectResponse("/target", headers=h) + assert resp.headers.get("Location") == "/target" + assert resp.headers.get("X-Custom") == "value" + + +def test_redirect_response_invalid_status_code(): + with pytest.raises(ValueError, match="Invalid redirect status code"): + RedirectResponse("/target", status_code=200)