fix: file serving (sparckles#881)

* fix: file serving
tavallaie · Jul 3, 2024 · 0e2910c · 0e2910c
1 parent 1e33769
commit 0e2910c
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 12 deletions.
diff --git a/docs_src/src/pages/documentation/api_reference/file-uploads.mdx b/docs_src/src/pages/documentation/api_reference/file-uploads.mdx
@@ -189,7 +189,7 @@ Batman scaled his application across multiple cores for better performance. He u
 
     @app.get("/")
     async def h(request):
-        return serve_file("./index.html")
+        return serve_file("./index.html", file_name="index.html") # file_name is optional
 
     app.start(port=8080)
 
@@ -203,7 +203,7 @@ Batman scaled his application across multiple cores for better performance. He u
 
     @app.get("/")
     async def h(request: Request):
-        return serve_file("./index.html")
+        return serve_file("./index.html", file_name="index.html") # file_name is optional
 
     app.start(port=8080)
 

diff --git a/integration_tests/test_file_download.py b/integration_tests/test_file_download.py
@@ -6,6 +6,6 @@
 @pytest.mark.parametrize("function_type", ["sync", "async"])
 def test_file_download(function_type: str, session):
     r = get(f"/{function_type}/file/download")
-    assert r.headers.get("Content-Disposition") == "attachment"
+    assert r.headers.get("Content-Disposition") == "attachment; filename=test.txt"
 
     assert r.text == "This is a test file for the downloading purpose"
diff --git a/robyn/responses.py b/robyn/responses.py
@@ -1,4 +1,6 @@
 from typing import Optional
+import os
+import mimetypes
 
 from robyn.robyn import Response, Headers
 
@@ -39,11 +41,21 @@ def serve_html(file_path: str) -> FileResponse:
     return FileResponse(file_path, headers=Headers({"Content-Type": "text/html"}))
 
 
-def serve_file(file_path: str) -> FileResponse:
+def serve_file(file_path: str, file_name: Optional[str] = None) -> FileResponse:
     """
     This function will help in serving a file
 
     :param file_path str: file path to serve as a response
+    :param file_name [str | None]: file name to serve as a response, defaults to None
     """
+    file_name = file_name or os.path.basename(file_path)
 
-    return FileResponse(file_path, headers=Headers({"Content-Disposition": "attachment"}))
+    mime_type = mimetypes.guess_type(file_name)[0]
+
+    headers = Headers({"Content-Type": mime_type})
+    headers.append("Content-Disposition", f"attachment; filename={file_name}")
+
+    return FileResponse(
+        file_path,
+        headers=headers,
+    )
diff --git a/src/io_helpers/mod.rs b/src/io_helpers/mod.rs
@@ -25,9 +25,9 @@ pub fn apply_hashmap_headers(response: &mut HttpResponseBuilder, headers: &Heade
 /// * `file_path` - The file path that we want the function to read
 ///
 // ideally this should be async
-pub fn read_file(file_path: &str) -> Result<String> {
+pub fn read_file(file_path: &str) -> Result<Vec<u8>> {
     let mut file = File::open(file_path)?;
     let mut buf = vec![];
     file.read_to_end(&mut buf)?;
-    Ok(String::from_utf8_lossy(&buf).to_string())
+    Ok(buf)
 }
diff --git a/src/types/response.rs b/src/types/response.rs
@@ -148,12 +148,15 @@ impl PyResponse {
 
     #[setter]
     pub fn set_file_path(&mut self, py: Python, file_path: &str) -> PyResult<()> {
-        // we should be handling based on headers but works for now
         self.response_type = "static_file".to_string();
         self.file_path = Some(file_path.to_string());
-        self.description = read_file(file_path)
-            .map_err(|e| PyErr::new::<PyIOError, _>(e.to_string()))?
-            .into_py(py);
-        Ok(())
+
+        match read_file(file_path) {
+            Ok(content) => {
+                self.description = PyBytes::new(py, &content).into();
+                Ok(())
+            }
+            Err(e) => Err(PyIOError::new_err(format!("Failed to read file: {}", e))),
+        }
     }
 }