Merge pull request RustPython#1226 from palaviv/subprocess-2

Extend subprocess.Popen

Author: windelbouwman

tests/snippets/stdlib_subprocess.py

@@ -1,6 +1,7 @@
 import subprocess
 import time
 import sys
+import signal
 
 from testutils import assertRaises
 
@@ -29,7 +30,29 @@
 
 if "win" not in sys.platform:
 	# unix
-	assert p.stdout.read() == b"test\n"
+	test_output = b"test\n"
 else:
 	# windows
-	assert p.stdout.read() == b"test\r\n"
+	test_output = b"test\r\n"
+
+assert p.stdout.read() == test_output
+
+p = subprocess.Popen(["sleep", "2"])
+p.terminate()
+p.wait()
+if "win" not in sys.platform:
+	assert p.returncode == -signal.SIGTERM
+else:
+	assert p.returncode == 1
+
+p = subprocess.Popen(["sleep", "2"])
+p.kill()
+p.wait()
+if "win" not in sys.platform:
+	assert p.returncode == -signal.SIGKILL
+else:
+	assert p.returncode == 1
+
+p = subprocess.Popen(["echo", "test"], stdout=subprocess.PIPE)
+(stdout, stderr) = p.communicate()
+assert stdout == test_output

vm/src/obj/objbytes.rs

@@ -11,7 +11,8 @@ use std::ops::Deref;
 
 use crate::function::OptionalArg;
 use crate::pyobject::{
-    PyClassImpl, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
+    IntoPyObject, PyClassImpl, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue,
+    TryFromObject,
 };
 
 use super::objbyteinner::{
@@ -59,6 +60,12 @@ impl PyBytes {
     }
 }
 
+impl IntoPyObject for Vec<u8> {
+    fn into_pyobject(self, vm: &VirtualMachine) -> PyResult {
+        Ok(vm.ctx.new_bytes(self))
+    }
+}
+
 impl Deref for PyBytes {
     type Target = [u8];
 

vm/src/obj/objtuple.rs

@@ -3,7 +3,9 @@ use std::fmt;
 
 use crate::function::OptionalArg;
 use crate::pyhash;
-use crate::pyobject::{IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
+use crate::pyobject::{
+    IdProtocol, IntoPyObject, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue,
+};
 use crate::vm::{ReprGuard, VirtualMachine};
 
 use super::objbool;
@@ -37,6 +39,27 @@ impl PyValue for PyTuple {
     }
 }
 
+impl<A> IntoPyObject for (A,)
+where
+    A: IntoPyObject,
+{
+    fn into_pyobject(self, vm: &VirtualMachine) -> PyResult {
+        Ok(vm.ctx.new_tuple(vec![self.0.into_pyobject(vm)?]))
+    }
+}
+
+impl<A, B> IntoPyObject for (A, B)
+where
+    A: IntoPyObject,
+    B: IntoPyObject,
+{
+    fn into_pyobject(self, vm: &VirtualMachine) -> PyResult {
+        Ok(vm
+            .ctx
+            .new_tuple(vec![self.0.into_pyobject(vm)?, self.1.into_pyobject(vm)?]))
+    }
+}
+
 impl PyTuple {
     pub fn fast_getitem(&self, idx: usize) -> PyObjectRef {
         self.elements[idx].clone()

vm/src/stdlib/os.rs

@@ -159,7 +159,7 @@ pub fn os_open(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.ctx.new_int(raw_file_number(handle)))
 }
 
-fn convert_io_error(vm: &VirtualMachine, err: io::Error) -> PyObjectRef {
+pub fn convert_io_error(vm: &VirtualMachine, err: io::Error) -> PyObjectRef {
     let os_error = match err.kind() {
         ErrorKind::NotFound => {
             let exc_type = vm.ctx.exceptions.file_not_found_error.clone();

vm/src/stdlib/subprocess.rs

@@ -1,17 +1,19 @@
 use std::cell::RefCell;
+use std::ffi::OsString;
 use std::fs::File;
 use std::time::Duration;
 
 use subprocess;
 
 use crate::function::OptionalArg;
+use crate::obj::objbytes::PyBytesRef;
 use crate::obj::objlist::PyListRef;
 use crate::obj::objsequence;
 use crate::obj::objstr::{self, PyStringRef};
 use crate::obj::objtype::PyClassRef;
 use crate::pyobject::{Either, IntoPyObject, PyObjectRef, PyRef, PyResult, PyValue};
 use crate::stdlib::io::io_open;
-use crate::stdlib::os::{raw_file_number, rust_file};
+use crate::stdlib::os::{convert_io_error, raw_file_number, rust_file};
 use crate::vm::VirtualMachine;
 
 #[derive(Debug)]
@@ -37,6 +39,8 @@ struct PopenArgs {
     stdout: Option<i64>,
     #[pyarg(positional_or_keyword, default = "None")]
     stderr: Option<i64>,
+    #[pyarg(positional_or_keyword, default = "None")]
+    cwd: Option<PyStringRef>,
 }
 
 impl IntoPyObject for subprocess::ExitStatus {
@@ -95,13 +99,15 @@ impl PopenRef {
                 .map(|x| objstr::get_value(x))
                 .collect(),
         };
+        let cwd = args.cwd.map(|x| OsString::from(x.as_str()));
 
         let process = subprocess::Popen::create(
             &command_list,
             subprocess::PopenConfig {
                 stdin,
                 stdout,
                 stderr,
+                cwd,
                 ..Default::default()
             },
         )
@@ -149,6 +155,36 @@ impl PopenRef {
     fn stderr(self, vm: &VirtualMachine) -> PyResult {
         convert_to_file_io(&self.process.borrow().stderr, "rb".to_string(), vm)
     }
+
+    fn terminate(self, vm: &VirtualMachine) -> PyResult<()> {
+        self.process
+            .borrow_mut()
+            .terminate()
+            .map_err(|err| convert_io_error(vm, err))
+    }
+
+    fn kill(self, vm: &VirtualMachine) -> PyResult<()> {
+        self.process
+            .borrow_mut()
+            .kill()
+            .map_err(|err| convert_io_error(vm, err))
+    }
+
+    #[allow(clippy::type_complexity)]
+    fn communicate(
+        self,
+        stdin: OptionalArg<PyBytesRef>,
+        vm: &VirtualMachine,
+    ) -> PyResult<(Option<Vec<u8>>, Option<Vec<u8>>)> {
+        self.process
+            .borrow_mut()
+            .communicate_bytes(stdin.into_option().as_ref().map(|bytes| bytes.get_value()))
+            .map_err(|err| convert_io_error(vm, err))
+    }
+
+    fn pid(self, _vm: &VirtualMachine) -> Option<u32> {
+        self.process.borrow().pid()
+    }
 }
 
 pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
@@ -165,6 +201,10 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
         "stdin" => ctx.new_property(PopenRef::stdin),
         "stdout" => ctx.new_property(PopenRef::stdout),
         "stderr" => ctx.new_property(PopenRef::stderr),
+        "terminate" => ctx.new_rustfunc(PopenRef::terminate),
+        "kill" => ctx.new_rustfunc(PopenRef::kill),
+        "communicate" => ctx.new_rustfunc(PopenRef::communicate),
+        "pid" => ctx.new_property(PopenRef::pid),
     });
 
     let module = py_module!(vm, "subprocess", {