Merge pull request #2851 from deantvv/fix-os-sendfile

coolreader18 · web-flow · commit cc40ae025ec4 · 2021-08-13T20:34:21.000-05:00
Fix os sendfile
diff --git a/Lib/test/test_os.py b/Lib/test/test_os.py
@@ -3050,7 +3050,6 @@ def handle_error(self):
         raise
 
 
-@unittest.skip("TODO: RUSTPYTHON")
 @unittest.skipUnless(hasattr(os, 'sendfile'), "test needs os.sendfile()")
 class TestSendfile(unittest.TestCase):
 
@@ -3178,11 +3177,12 @@ def test_invalid_offset(self):
 
     def test_keywords(self):
         # Keyword arguments should be supported
-        os.sendfile(out=self.sockno, offset=0, count=4096,
-            **{'in': self.fileno})
+        os.sendfile(out_fd=self.sockno, in_fd=self.fileno,
+                    offset=0, count=4096)
         if self.SUPPORT_HEADERS_TRAILERS:
-            os.sendfile(self.sockno, self.fileno, offset=0, count=4096,
-                headers=(), trailers=(), flags=0)
+            os.sendfile(out_fd=self.sockno, in_fd=self.fileno,
+                        offset=0, count=4096,
+                        headers=(), trailers=(), flags=0)
 
     # --- headers / trailers tests
 
diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs
@@ -498,34 +498,73 @@ mod _os {
         fd.map(|fd| fd.0).map_err(|e| e.into_pyexception(vm))
     }
 
+    #[cfg(any(target_os = "linux", target_os = "macos"))]
+    #[derive(FromArgs)]
+    struct SendFileArgs {
+        #[pyarg(any)]
+        out_fd: i32,
+        #[pyarg(any)]
+        in_fd: i32,
+        #[pyarg(any)]
+        offset: i64,
+        #[pyarg(any)]
+        count: i64,
+        #[cfg(target_os = "macos")]
+        #[pyarg(any, optional)]
+        headers: OptionalArg<PyObjectRef>,
+        #[cfg(target_os = "macos")]
+        #[pyarg(any, optional)]
+        trailers: OptionalArg<PyObjectRef>,
+        #[cfg(target_os = "macos")]
+        #[allow(dead_code)]
+        #[pyarg(any, default)]
+        // TODO: not implemented
+        flags: OptionalArg<i32>,
+    }
+
     #[cfg(target_os = "linux")]
     #[pyfunction]
-    fn sendfile(out_fd: i32, in_fd: i32, offset: i64, count: u64, vm: &VirtualMachine) -> PyResult {
-        let mut file_offset = offset;
+    fn sendfile(args: SendFileArgs, vm: &VirtualMachine) -> PyResult {
+        let mut file_offset = args.offset;
 
-        let res =
-            nix::sys::sendfile::sendfile(out_fd, in_fd, Some(&mut file_offset), count as usize)
-                .map_err(|err| err.into_pyexception(vm))?;
+        let res = nix::sys::sendfile::sendfile(
+            args.out_fd,
+            args.in_fd,
+            Some(&mut file_offset),
+            args.count as usize,
+        )
+        .map_err(|err| err.into_pyexception(vm))?;
         Ok(vm.ctx.new_int(res as u64))
     }
 
     #[cfg(target_os = "macos")]
-    #[pyfunction]
-    #[allow(clippy::too_many_arguments)]
-    fn sendfile(
-        out_fd: i32,
-        in_fd: i32,
-        offset: i64,
-        count: i64,
-        headers: OptionalArg<PyObjectRef>,
-        trailers: OptionalArg<PyObjectRef>,
-        _flags: OptionalArg<i32>,
+    fn _extract_vec_bytes(
+        x: OptionalArg,
         vm: &VirtualMachine,
-    ) -> PyResult {
-        let headers = match headers.into_option() {
-            Some(x) => Some(vm.extract_elements::<PyBytesLike>(&x)?),
+    ) -> PyResult<Option<Vec<PyBytesLike>>> {
+        let inner = match x.into_option() {
+            Some(v) => {
+                let v = vm.extract_elements::<PyBytesLike>(&v)?;
+                if v.is_empty() {
+                    None
+                } else {
+                    Some(v)
+                }
+            }
             None => None,
         };
+        Ok(inner)
+    }
+
+    #[cfg(target_os = "macos")]
+    #[pyfunction]
+    fn sendfile(args: SendFileArgs, vm: &VirtualMachine) -> PyResult {
+        let headers = _extract_vec_bytes(args.headers, vm)?;
+        let count = headers
+            .as_ref()
+            .map(|v| v.iter().map(|s| s.len()).sum())
+            .unwrap_or(0) as i64
+            + args.count;
 
         let headers = headers
             .as_ref()
@@ -535,10 +574,7 @@ mod _os {
             .map(|v| v.iter().map(|borrowed| &**borrowed).collect::<Vec<_>>());
         let headers = headers.as_deref();
 
-        let trailers = match trailers.into_option() {
-            Some(x) => Some(vm.extract_elements::<PyBytesLike>(&x)?),
-            None => None,
-        };
+        let trailers = _extract_vec_bytes(args.trailers, vm)?;
 
         let trailers = trailers
             .as_ref()
@@ -548,8 +584,14 @@ mod _os {
             .map(|v| v.iter().map(|borrowed| &**borrowed).collect::<Vec<_>>());
         let trailers = trailers.as_deref();
 
-        let (res, written) =
-            nix::sys::sendfile::sendfile(in_fd, out_fd, offset, Some(count), headers, trailers);
+        let (res, written) = nix::sys::sendfile::sendfile(
+            args.in_fd,
+            args.out_fd,
+            args.offset,
+            Some(count),
+            headers,
+            trailers,
+        );
         res.map_err(|err| err.into_pyexception(vm))?;
         Ok(vm.ctx.new_int(written as u64))
     }
