use pyo3::{exceptions::PyTypeError, prelude::*};
use pyo3::types::{PyBytes, PyString, PyType};
use std::io;
use std::io::{Read, Seek, SeekFrom, Write};
#[derive(Debug)]
pub struct PyFileLikeObject {
inner: PyObject,
is_text_io: bool,
}
impl PyFileLikeObject {
pub fn new(object: PyObject) -> PyResult<Self> {
Python::with_gil(|py| {
let io = PyModule::import(py, "io")?;
let text_io = io.getattr("TextIOBase")?;
let text_io_type = text_io.extract::<&PyType>()?;
let is_text_io = object.as_ref(py).is_instance(text_io_type)?;
Ok(PyFileLikeObject {
inner: object,
is_text_io,
})
})
}
pub fn with_requirements(
object: PyObject,
read: bool,
write: bool,
seek: bool,
) -> PyResult<Self> {
Python::with_gil(|py| {
if read && object.getattr(py, "read").is_err() {
return Err(PyErr::new::<PyTypeError, _>(
"Object does not have a .read() method.",
));
}
if seek && object.getattr(py, "seek").is_err() {
return Err(PyErr::new::<PyTypeError, _>(
"Object does not have a .seek() method.",
));
}
if write && object.getattr(py, "write").is_err() {
return Err(PyErr::new::<PyTypeError, _>(
"Object does not have a .write() method.",
));
}
PyFileLikeObject::new(object)
})
}
}
fn pyerr_to_io_err(e: PyErr) -> io::Error {
Python::with_gil(|py| {
let e_as_object: PyObject = e.into_py(py);
match e_as_object.call_method(py, "__str__", (), None) {
Ok(repr) => match repr.extract::<String>(py) {
Ok(s) => io::Error::new(io::ErrorKind::Other, s),
Err(_e) => io::Error::new(io::ErrorKind::Other, "An unknown error has occurred"),
},
Err(_) => io::Error::new(io::ErrorKind::Other, "Err doesn't have __str__"),
}
})
}
impl Read for PyFileLikeObject {
fn read(&mut self, mut buf: &mut [u8]) -> Result<usize, io::Error> {
Python::with_gil(|py| {
if self.is_text_io {
if buf.len() < 4 {
return Err(io::Error::new(
io::ErrorKind::Other,
"buffer size must be at least 4 bytes",
));
}
let res = self
.inner
.call_method(py, "read", (buf.len() / 4,), None)
.map_err(pyerr_to_io_err)?;
let pystring: &PyString = res
.downcast(py)
.expect("Expecting to be able to downcast into str from read result.");
let bytes = pystring.to_str().unwrap().as_bytes();
buf.write_all(bytes)?;
Ok(bytes.len())
} else {
let res = self
.inner
.call_method(py, "read", (buf.len(),), None)
.map_err(pyerr_to_io_err)?;
let pybytes: &PyBytes = res
.downcast(py)
.expect("Expecting to be able to downcast into bytes from read result.");
let bytes = pybytes.as_bytes();
buf.write_all(bytes)?;
Ok(bytes.len())
}
})
}
}
impl Write for PyFileLikeObject {
fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
Python::with_gil(|py| {
let arg = if self.is_text_io {
let s = std::str::from_utf8(buf)
.expect("Tried to write non-utf8 data to a TextIO object.");
PyString::new(py, s).to_object(py)
} else {
PyBytes::new(py, buf).to_object(py)
};
let number_bytes_written = self
.inner
.call_method(py, "write", (arg,), None)
.map_err(pyerr_to_io_err)?;
if number_bytes_written.is_none(py) {
return Err(io::Error::new(
io::ErrorKind::Other,
"write() returned None, expected number of bytes written",
));
}
number_bytes_written.extract(py).map_err(pyerr_to_io_err)
})
}
fn flush(&mut self) -> Result<(), io::Error> {
Python::with_gil(|py| {
self.inner
.call_method(py, "flush", (), None)
.map_err(pyerr_to_io_err)?;
Ok(())
})
}
}
impl Seek for PyFileLikeObject {
fn seek(&mut self, pos: SeekFrom) -> Result<u64, io::Error> {
Python::with_gil(|py| {
let (whence, offset) = match pos {
SeekFrom::Start(i) => (0, i as i64),
SeekFrom::Current(i) => (1, i),
SeekFrom::End(i) => (2, i),
};
let new_position = self
.inner
.call_method(py, "seek", (offset, whence), None)
.map_err(pyerr_to_io_err)?;
new_position.extract(py).map_err(pyerr_to_io_err)
})
}
}