From c02527a6b672dd83c995311ba59438ae4d36d927 Mon Sep 17 00:00:00 2001
From: Fabian Freyer <fabian.freyer@physik.tu-berlin.de>
Date: Wed, 2 Jan 2019 19:52:03 +0100
Subject: [PATCH] python: add stdio methods to child

---
 bindings/python/src/child.rs | 106 +++++++++++++++++++++++++++++++++++
 1 file changed, 106 insertions(+)

diff --git a/bindings/python/src/child.rs b/bindings/python/src/child.rs
index 909ab5147..b8cc417aa 100644
--- a/bindings/python/src/child.rs
+++ b/bindings/python/src/child.rs
@@ -1,7 +1,10 @@
 use pyo3::prelude::*;
+use pyo3::types::PyByteArray;
 use pyo3::{exceptions, PyObjectWithToken};
 
+use std::io::{Read, Write};
 use std::ops::{Deref, DerefMut};
+use std::os::unix::io::AsRawFd;
 use std::os::unix::process::ExitStatusExt;
 
 #[pyclass]
@@ -92,6 +95,24 @@ impl Child {
     pub fn create(token: PyToken, inner: std::process::Child) -> Self {
         Child { token, inner }
     }
+
+    fn stdin(&mut self) -> PyResult<&mut std::process::ChildStdin> {
+        self.stdin
+            .as_mut()
+            .ok_or(exceptions::IOError::py_err("Stdin not captured"))
+    }
+
+    fn stdout(&mut self) -> PyResult<&mut std::process::ChildStdout> {
+        self.stdout
+            .as_mut()
+            .ok_or(exceptions::IOError::py_err("Stdout not captured"))
+    }
+
+    fn stderr(&mut self) -> PyResult<&mut std::process::ChildStderr> {
+        self.stderr
+            .as_mut()
+            .ok_or(exceptions::IOError::py_err("Stderr not captured"))
+    }
 }
 
 #[pymethods]
@@ -126,6 +147,11 @@ impl Child {
         })
     }
 
+    #[getter]
+    fn get_stdin_fd(&mut self) -> PyResult<std::os::unix::io::RawFd> {
+        Ok(self.stdin()?.as_raw_fd())
+    }
+
     fn try_wait(&mut self) -> PyResult<Option<Py<ExitStatus>>> {
         self.inner
             .try_wait()
@@ -133,4 +159,84 @@ impl Child {
             .map(|s| self.py().init(|token| ExitStatus { token, inner: s }))
             .map_or(Ok(None), |v| v.map(Some))
     }
+
+    pub fn write_stdin(&mut self, buf: &PyByteArray) -> PyResult<usize> {
+        self.stdin()?
+            .write(buf.data())
+            .map_err(|_| exceptions::IOError::py_err("Could not write to stdin"))
+    }
+
+    pub fn write_stdin_str(&mut self, buf: String) -> PyResult<()> {
+        self.stdin()?
+            .write_all(buf.as_bytes())
+            .map_err(|_| exceptions::IOError::py_err("Could not write to stdin"))
+    }
+
+    pub fn flush_stdin(&mut self) -> PyResult<()> {
+        self.stdin()?
+            .flush()
+            .map_err(|e| PyErr::new::<exceptions::IOError, String>(format!("{}", e)))
+    }
+
+    #[getter]
+    fn get_stdout_fd(&mut self) -> PyResult<std::os::unix::io::RawFd> {
+        Ok(self.stdout()?.as_raw_fd())
+    }
+
+    pub fn read_stdout(&mut self, len: usize) -> PyResult<&PyByteArray> {
+        let into: PyResult<Vec<u8>> = {
+            let stdout = self.stdout()?;
+
+            let mut into = vec![0; len];
+            let read = stdout
+                .read(&mut into)
+                .map_err(|_| exceptions::IOError::py_err("Could not read from Stdout"))?;
+            into.truncate(read);
+            Ok(into)
+        };
+
+        Ok(PyByteArray::new(self.py(), &into?))
+    }
+
+    pub fn readall_stdout_str(&mut self) -> PyResult<String> {
+        let stdout = self.stdout()?;
+
+        let mut into = String::new();
+        stdout
+            .read_to_string(&mut into)
+            .map_err(|_| exceptions::IOError::py_err("Could not read from Stdout"))?;
+
+        Ok(into)
+    }
+
+    #[getter]
+    fn get_stderr_fd(&mut self) -> PyResult<std::os::unix::io::RawFd> {
+        Ok(self.stderr()?.as_raw_fd())
+    }
+
+    pub fn read_stderr(&mut self, len: usize) -> PyResult<&PyByteArray> {
+        let into: PyResult<Vec<u8>> = {
+            let stderr = self.stderr()?;
+
+            let mut into = vec![0; len];
+            let read = stderr
+                .read(&mut into)
+                .map_err(|_| exceptions::IOError::py_err("Could not read from Stderr"))?;
+            into.truncate(read);
+            Ok(into)
+        };
+
+        Ok(PyByteArray::new(self.py(), &into?))
+    }
+
+    pub fn readall_stderr_str(&mut self) -> PyResult<String> {
+        let stderr = self.stderr()?;
+
+        let mut into = String::new();
+        stderr
+            .read_to_string(&mut into)
+            .map_err(|_| exceptions::IOError::py_err("Could not read from Stdout"))?;
+
+        Ok(into)
+    }
 }