Changed Response to use body: bytes (#375)

* Changed Response to use body: bytes

- Added ActixBytesWrapper to conform to pyclass and MessageBody

* Made suggested improvements to allow both PyString and PyBytes

- Reverted unnecessary casting and type changes

* Added test for supplying unsupported type to body

- WIP: Issue with async testing

* Remove unused unsupported type test due to issues with raises

* Fixed issue where .unwrap causes Panic and instead raises Exception

- Added tests for bad and good body types

* Removed unused return

- Parameterized body test fixtures

---------

Co-authored-by: Antoine Romero-Romero <45259848+AntoineRR@users.noreply.github.com>

Author: madhavajay

integration_tests/base_routes.py

@@ -444,7 +444,32 @@ async def async_body_patch(request):
     return bytearray(request["body"]).decode("utf-8")
 
 
-# ===== Main =====
+@app.get("/binary_output_sync")
+def binary_output_sync(request):
+    return b"OK"
+
+
+@app.get("/binary_output_response_sync")
+def binary_output_response_sync(request):
+    return Response(
+        status_code=200,
+        headers={"Content-Type": "application/octet-stream"},
+        body="OK",
+    )
+
+
+@app.get("/binary_output_async")
+async def binary_output_async(request):
+    return b"OK"
+
+
+@app.get("/binary_output_response_async")
+async def binary_output_response_async(request):
+    return Response(
+        status_code=200,
+        headers={"Content-Type": "application/octet-stream"},
+        body="OK",
+    )
 
 
 @app.get("/sync/raise")
@@ -457,6 +482,9 @@ async def async_raise():
     raise Exception()
 
 
+# ===== Main =====
+
+
 if __name__ == "__main__":
     app.add_request_header("server", "robyn")
     app.add_directory(

integration_tests/test_binary_output.py

@@ -0,0 +1,31 @@
+import requests
+
+BASE_URL = "http://127.0.0.1:8080"
+
+
+def test_binary_output_sync(session):
+    r = requests.get(f"{BASE_URL}/binary_output_sync")
+    assert r.status_code == 200
+    assert r.headers["Content-Type"] == "application/octet-stream"
+    assert r.text == "OK"
+
+
+def test_binary_output_response_sync(session):
+    r = requests.get(f"{BASE_URL}/binary_output_response_sync")
+    assert r.status_code == 200
+    assert r.headers["Content-Type"] == "application/octet-stream"
+    assert r.text == "OK"
+
+
+def test_binary_output_async(session):
+    r = requests.get(f"{BASE_URL}/binary_output_async")
+    assert r.status_code == 200
+    assert r.headers["Content-Type"] == "application/octet-stream"
+    assert r.text == "OK"
+
+
+def test_binary_output_response_async(session):
+    r = requests.get(f"{BASE_URL}/binary_output_response_async")
+    assert r.status_code == 200
+    assert r.headers["Content-Type"] == "application/octet-stream"
+    assert r.text == "OK"

integration_tests/test_unsupported_types.py

@@ -0,0 +1,43 @@
+import pytest
+
+from robyn.robyn import Response
+
+
+class A:
+    pass
+
+
+bad_bodies = [
+    None,
+    1,
+    True,
+    A,
+    {"body": "OK"},
+    ["OK", b"OK"],
+    Response(
+        status_code=200,
+        headers={},
+        body=b"OK",
+    ),
+]
+
+good_bodies = ["OK", b"OK"]
+
+
+@pytest.mark.parametrize("body", bad_bodies)
+def test_bad_body_types(body):
+    with pytest.raises(ValueError):
+        _ = Response(
+            status_code=200,
+            headers={},
+            body=body,
+        )
+
+
+@pytest.mark.parametrize("body", good_bodies)
+def test_good_body_types(body):
+    _ = Response(
+        status_code=200,
+        headers={},
+        body=body,
+    )

robyn/robyn.pyi

@@ -19,7 +19,7 @@ class FunctionInfo:
 class Response:
     status_code: int
     headers: dict[str, str]
-    body: str
+    body: bytes
 
     def set_file_path(self, file_path: str):
         pass

robyn/router.py

@@ -39,11 +39,18 @@ def _format_response(self, res):
                 response.set_file_path(file_path)
         elif type(res) == Response:
             response = res
+        elif type(res) == bytes:
+            response = Response(
+                status_code=200,
+                headers={"Content-Type": "application/octet-stream"},
+                body=res,
+            )
         else:
             response = Response(
-                status_code=200, headers={"Content-Type": "text/plain"}, body=str(res)
+                status_code=200,
+                headers={"Content-Type": "text/plain"},
+                body=str(res).encode("utf-8"),
             )
-
         return response
 
     def add_route(

src/types.rs

@@ -1,13 +1,89 @@
+use core::{mem};
+use std::ops::{Deref, DerefMut};
+use std::{
+    convert::Infallible,
+    task::{Context, Poll},
+    pin::Pin,
+};
 use std::collections::HashMap;
 
 use actix_web::web::Bytes;
 use actix_web::{http::Method, HttpRequest};
+use actix_http::body::MessageBody;
+use actix_http::body::BodySize;
+
 use anyhow::Result;
 use dashmap::DashMap;
 use pyo3::{exceptions, prelude::*};
+use pyo3::types::{PyBytes, PyString};
+use pyo3::exceptions::PyValueError;
 
 use crate::io_helpers::read_file;
 
+fn type_of<T>(_: &T) -> String {
+    std::any::type_name::<T>().to_string()
+}
+
+#[derive(Debug, Clone)]
+#[pyclass]
+pub struct ActixBytesWrapper(Bytes);
+
+// provides an interface between pyo3::types::{PyString, PyBytes} and actix_web::web::Bytes
+impl ActixBytesWrapper {
+    pub fn new(raw_body: &PyAny) -> PyResult<Self> {
+        let value = if let Ok(v) = raw_body.downcast::<PyString>() {
+            v.to_string().into_bytes()
+        } else if let Ok(v) = raw_body.downcast::<PyBytes>() {
+            v.as_bytes().to_vec()
+        } else {
+            return Err(PyValueError::new_err(
+                format!("Could not convert {} specified body to bytes", type_of(raw_body))
+            ));
+        };
+        Ok(Self(Bytes::from(value)))
+    }
+}
+
+impl Deref for ActixBytesWrapper {
+    type Target = Bytes;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for ActixBytesWrapper {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl MessageBody for ActixBytesWrapper {
+    type Error = Infallible;
+
+    #[inline]
+    fn size(&self) -> BodySize {
+        BodySize::Sized(self.len() as u64)
+    }
+
+    #[inline]
+    fn poll_next(
+        self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Bytes, Self::Error>>> {
+        if self.is_empty() {
+            Poll::Ready(None)
+        } else {
+            Poll::Ready(Some(Ok(mem::take(self.get_mut()))))
+        }
+    }
+
+    #[inline]
+    fn try_into_bytes(self) -> Result<Bytes, Self> {
+        Ok(self.0)
+    }
+}
+
 #[pyclass]
 #[derive(Debug, Clone)]
 pub struct FunctionInfo {
@@ -81,32 +157,33 @@ pub struct Response {
     pub status_code: u16,
     pub response_type: String,
     pub headers: HashMap<String, String>,
-    pub body: String,
+    pub body: ActixBytesWrapper,
     pub file_path: Option<String>,
 }
 
 #[pymethods]
 impl Response {
     #[new]
-    pub fn new(status_code: u16, headers: HashMap<String, String>, body: String) -> Self {
-        Self {
+    pub fn new(status_code: u16, headers: HashMap<String, String>, body: &PyAny) -> PyResult<Self> {
+        Ok(Self {
             status_code,
             // we should be handling based on headers but works for now
             response_type: "text".to_string(),
             headers,
-            body,
+            body: ActixBytesWrapper::new(body)?,
             file_path: None,
-        }
+        })
     }
 
     pub fn set_file_path(&mut self, 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.body = match read_file(file_path) {
+        let response = match read_file(file_path) {
             Ok(b) => b,
             Err(e) => return Err(exceptions::PyIOError::new_err::<String>(e.to_string())),
         };
+        self.body = ActixBytesWrapper(Bytes::from(response));
         Ok(())
     }
 }