Skip to content

Commit

Permalink
python: add stdio methods to child
Browse files Browse the repository at this point in the history
  • Loading branch information
fabianfreyer committed Jan 2, 2019
1 parent 62d5cc0 commit c02527a
Showing 1 changed file with 106 additions and 0 deletions.
106 changes: 106 additions & 0 deletions bindings/python/src/child.rs
Original file line number Diff line number Diff line change
@@ -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]
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -126,11 +147,96 @@ 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()
.map_err(|e| PyErr::new::<exceptions::IOError, String>(format!("{}", e)))?
.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)
}
}

0 comments on commit c02527a

Please sign in to comment.