Minor enhancements on multipart files handling

Author: gi0baro

emmett_core/http/wrappers/helpers.py

@@ -56,6 +56,9 @@ class FileStorage:
     def __init__(self, file):
         self.file = file
 
+    def __iter__(self):
+        return self.file.__iter__()
+
     def __getattr__(self, name):
         return getattr(self.file, name)
 

src/multipart/parse.rs

@@ -175,6 +175,7 @@ impl MultiPartParser {
             }
 
             if let MultiPartParserState::File(filepart) = &mut self.state {
+                // potentially allow py threads?
                 let (read, found) = reader.stream_until_token(
                     lt_boundary,
                     &mut filepart.file.as_mut().expect("uninitialized file part"),
@@ -189,6 +190,8 @@ impl MultiPartParser {
                 let state = mem::take(&mut self.state);
                 match state {
                     MultiPartParserState::File(part) => {
+                        // potentially allow py threads?
+                        part.file.as_ref().unwrap().sync_data()?;
                         self.stack.push_back(Node::File(part));
                     }
                     _ => unreachable!(),

src/multipart/parts.rs

@@ -3,7 +3,7 @@ use http::header::{self, HeaderMap};
 use pyo3::{exceptions::PyStopIteration, prelude::*, types::PyBytes};
 use std::{
     fs::File,
-    io::{BufRead, BufReader, Read},
+    io::{BufReader, Read},
     path::PathBuf,
     sync::Mutex,
 };
@@ -117,7 +117,7 @@ impl FilePartReader {
         drop(inner.file.take().expect("uninitialized file part"));
         let file = File::open(inner.path.clone()).map_err::<anyhow::Error, _>(|_| error_io!())?;
         let size = file.metadata().unwrap().len();
-        let reader = Mutex::new(BufReader::with_capacity(4096, file));
+        let reader = Mutex::new(BufReader::with_capacity(131_072, file));
         Ok(Self { inner, reader, size })
     }
 
@@ -128,7 +128,6 @@ impl FilePartReader {
         let mut len_read = 0;
 
         while len_read < size {
-            guard.fill_buf()?;
             let rsize = guard.read(&mut buf[len_read..])?;
             if rsize == 0 {
                 break;
@@ -143,8 +142,6 @@ impl FilePartReader {
 
     fn read_all(&self) -> Result<Vec<u8>> {
         let mut guard = self.reader.lock().unwrap();
-
-        guard.fill_buf()?;
         let mut buf = Vec::new();
         guard.read_to_end(&mut buf)?;
         Ok(buf)
@@ -174,8 +171,8 @@ impl FilePartReader {
     #[pyo3(signature = (size = None))]
     fn read<'p>(&self, py: Python<'p>, size: Option<usize>) -> Result<Bound<'p, PyBytes>> {
         let buf = match size {
-            Some(size) => self.read_chunk(size),
-            None => self.read_all(),
+            Some(size) => py.allow_threads(|| self.read_chunk(size)),
+            None => py.allow_threads(|| self.read_all()),
         }?;
         Ok(PyBytes::new(py, &buf[..]))
     }
@@ -185,7 +182,7 @@ impl FilePartReader {
     }
 
     fn __next__<'p>(&self, py: Python<'p>) -> Result<Bound<'p, PyBytes>> {
-        let buf = self.read_chunk(4096)?;
+        let buf = py.allow_threads(|| self.read_chunk(131_072))?;
         if buf.is_empty() {
             return Err(PyStopIteration::new_err(py.None()).into());
         }

tests/multipart/test_multipart.py

@@ -57,6 +57,37 @@ async def multipart():
     return client
 
 
+@pytest.fixture(scope="function")
+def multipart_stream_client(current, client, tmpdir: Path):
+    app = client.application
+
+    @app.route("/", output="str")
+    async def multipart():
+        target = tmpdir / "save.txt"
+        files = await current.request.files
+        with target.open("wb") as tf:
+            for chunk in files.test:
+                tf.write(chunk)
+        return ""
+
+    return client
+
+
+@pytest.fixture(scope="function")
+def multipart_copy_client(current, client, tmpdir: Path):
+    app = client.application
+
+    @app.route("/", output="str")
+    async def multipart():
+        target = tmpdir / "save.txt"
+        files = await current.request.files
+        with target.open("wb") as tf:
+            tf.write(files.test.read())
+        return ""
+
+    return client
+
+
 def test_multipart_request_data(multipart_client):
     response = multipart_client.post("/", data={"some": "data"}, content_type="multipart/form-data")
     assert response.json() == {"params": {"some": ["data"]}, "files": {}}
@@ -340,8 +371,8 @@ def test_multipart_request_file_save(tmpdir: Path, multipart_save_client):
     target = tmpdir / "save.txt"
     with path.open("wb") as file:
         file.write(b"<")
-        for i in range(8192):
-            file.write(f"{i}".zfill(5).encode("utf8"))
+        for i in range(8192 * 128):
+            file.write(f"{i}".zfill(7).encode("utf8"))
         file.write(b">")
 
     with path.open("rb") as f:
@@ -350,3 +381,37 @@ def test_multipart_request_file_save(tmpdir: Path, multipart_save_client):
 
     with path.open("rb") as f1, target.open("rb") as f2:
         assert f1.read() == f2.read()
+
+
+def test_multipart_request_file_stream(tmpdir: Path, multipart_stream_client):
+    path = tmpdir / "test.txt"
+    target = tmpdir / "save.txt"
+    with path.open("wb") as file:
+        file.write(b"<")
+        for i in range(8192 * 128):
+            file.write(f"{i}".zfill(7).encode("utf8"))
+        file.write(b">")
+
+    with path.open("rb") as f:
+        response = multipart_stream_client.post("/", data={"test": (f, "test.txt", "text/plain")})
+        assert response.status == 200
+
+    with path.open("rb") as f1, target.open("rb") as f2:
+        assert f1.read() == f2.read()
+
+
+def test_multipart_request_file_copy(tmpdir: Path, multipart_copy_client):
+    path = tmpdir / "test.txt"
+    target = tmpdir / "save.txt"
+    with path.open("wb") as file:
+        file.write(b"<")
+        for i in range(8192 * 128):
+            file.write(f"{i}".zfill(7).encode("utf8"))
+        file.write(b">")
+
+    with path.open("rb") as f:
+        response = multipart_copy_client.post("/", data={"test": (f, "test.txt", "text/plain")})
+        assert response.status == 200
+
+    with path.open("rb") as f1, target.open("rb") as f2:
+        assert f1.read() == f2.read()