Support timeout in Popen.communicate

palaviv · palaviv · commit 9c12b95cf497 · 2020-02-01T09:54:05.000+02:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/tests/snippets/stdlib_subprocess.py b/tests/snippets/stdlib_subprocess.py
@@ -51,3 +51,7 @@
 p = subprocess.Popen(["echo", "test"], stdout=subprocess.PIPE)
 (stdout, stderr) = p.communicate()
 assert stdout.strip() == b"test"
+
+p = subprocess.Popen(["sleep", "5"], stdout=subprocess.PIPE)
+with assert_raises(subprocess.TimeoutExpired):
+	p.communicate(timeout=1)
diff --git a/vm/Cargo.toml b/vm/Cargo.toml
@@ -85,7 +85,7 @@ adler32 = "1.0.3"
 flate2 = { version = "1.0", features = ["zlib"], default-features = false }
 libz-sys = "1.0.25"
 gethostname = "0.2.0"
-subprocess = "0.1.18"
+subprocess = "0.2.2"
 num_cpus = "1"
 socket2 = { version = "0.3", features = ["unix"] }
 
diff --git a/vm/src/stdlib/subprocess.rs b/vm/src/stdlib/subprocess.rs
@@ -1,6 +1,7 @@
 use std::cell::RefCell;
 use std::ffi::OsString;
 use std::fs::File;
+use std::io::ErrorKind;
 use std::time::Duration;
 
 use subprocess;
@@ -191,13 +192,21 @@ impl PopenRef {
         vm: &VirtualMachine,
     ) -> PyResult<(Option<Vec<u8>>, Option<Vec<u8>>)> {
         let bytes = match args.input {
-            OptionalArg::Present(ref bytes) => Some(bytes.get_value()),
+            OptionalArg::Present(ref bytes) => Some(bytes.get_value().to_vec()),
             OptionalArg::Missing => None,
         };
-        self.process
-            .borrow_mut()
-            .communicate_bytes(bytes)
-            .map_err(|err| convert_io_error(vm, err))
+        let mut communicator = self.process.borrow_mut().communicate_start(bytes);
+        if let OptionalArg::Present(timeout) = args.timeout {
+            communicator = communicator.limit_time(Duration::new(timeout, 0));
+        }
+        communicator.read().map_err(|err| {
+            if err.error.kind() == ErrorKind::TimedOut {
+                let timeout_expired = vm.try_class("_subprocess", "TimeoutExpired").unwrap();
+                vm.new_exception_msg(timeout_expired, "Timeout".to_string())
+            } else {
+                convert_io_error(vm, err.error)
+            }
+        })
     }
 
     fn pid(self, _vm: &VirtualMachine) -> Option<u32> {
