pyo3_file/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use pyo3::intern;
4use pyo3::{exceptions::PyTypeError, prelude::*};
5use std::borrow::Cow;
6
7use pyo3::types::{PyBytes, PyString};
8
9use std::io;
10use std::io::{Read, Seek, SeekFrom, Write};
11#[cfg(unix)]
12use std::os::fd::{AsRawFd, RawFd};
13
14/// A wrapper around a Python object that implements the file-like interface.
15#[derive(Debug)]
16pub struct PyFileLikeObject {
17    // We use PyObject instead of Bound<PyAny> because Bound<PyAny> is a GIL-bound type.
18    // We want to avoid holding the GIL when creating the struct.
19    // The GIL will be re-taken when the methods are called.
20    inner: PyObject,
21    is_text_io: bool,
22}
23
24impl Clone for PyFileLikeObject {
25    fn clone(&self) -> Self {
26        Python::with_gil(|py| PyFileLikeObject {
27            inner: self.inner.clone_ref(py),
28            is_text_io: self.is_text_io,
29        })
30    }
31}
32
33/// Wraps a `PyObject`, and implements read, seek, and write for it.
34impl PyFileLikeObject {
35    /// Creates an instance of a `PyFileLikeObject` from a `PyObject`.
36    ///
37    /// To assert the object has the required methods methods,
38    /// instantiate it with [`with_requirements`][Self::with_requirements].
39    ///
40    /// Prefer using [`py_new`][Self::py_new] if you already have a `Bound<PyAny>` object, as this
41    /// method re-acquires the GIL internally.
42    pub fn new(object: PyObject) -> PyResult<Self> {
43        Python::with_gil(|py| Self::py_new(object.into_bound(py)))
44    }
45
46    /// Same as `PyFileLikeObject::new`, but validates that the underlying
47    /// python object has a `read`, `write`, and `seek` methods in respect to parameters.
48    /// Will return a `TypeError` if object does not have `read`, `seek`, `write` and `fileno` methods.
49    ///
50    /// Prefer using [`py_with_requirements`][Self::py_with_requirements] if you already have a
51    /// `Bound<PyAny>` object, as this method re-acquires the GIL internally.
52    pub fn with_requirements(
53        object: PyObject,
54        read: bool,
55        write: bool,
56        seek: bool,
57        fileno: bool,
58    ) -> PyResult<Self> {
59        Python::with_gil(|py| {
60            Self::py_with_requirements(object.into_bound(py), read, write, seek, fileno)
61        })
62    }
63
64    /// Creates an instance of a `PyFileLikeObject` from a `PyObject`.
65    ///
66    /// Prefer using this instead of [`new`][Self::new] if you already have a `Bound<PyAny>`
67    /// object, as this method does not acquire the GIL internally.
68    pub fn py_new(obj: Bound<PyAny>) -> PyResult<Self> {
69        let text_io = consts::text_io_base(obj.py())?;
70        let is_text_io = obj.is_instance(text_io)?;
71
72        Ok(PyFileLikeObject {
73            inner: obj.unbind(),
74            is_text_io,
75        })
76    }
77
78    /// Creates an instance of a `PyFileLikeObject` from a `PyObject`.
79    ///
80    /// Prefer using this instead of [`with_requirements`][Self::with_requirements] if you already
81    /// have a `Bound<PyAny>` object, as this method does not acquire the GIL internally.
82    pub fn py_with_requirements(
83        obj: Bound<PyAny>,
84        read: bool,
85        write: bool,
86        seek: bool,
87        fileno: bool,
88    ) -> PyResult<Self> {
89        if read && !obj.hasattr(consts::read(obj.py()))? {
90            return Err(PyTypeError::new_err(
91                "Object does not have a .read() method.",
92            ));
93        }
94
95        if seek && !obj.hasattr(consts::seek(obj.py()))? {
96            return Err(PyTypeError::new_err(
97                "Object does not have a .seek() method.",
98            ));
99        }
100
101        if write && !obj.hasattr(consts::write(obj.py()))? {
102            return Err(PyTypeError::new_err(
103                "Object does not have a .write() method.",
104            ));
105        }
106
107        if fileno && !obj.hasattr(consts::fileno(obj.py()))? {
108            return Err(PyTypeError::new_err(
109                "Object does not have a .fileno() method.",
110            ));
111        }
112
113        PyFileLikeObject::py_new(obj)
114    }
115
116    pub fn py_read(&self, py: Python<'_>, mut buf: &mut [u8]) -> io::Result<usize> {
117        let inner = self.inner.bind(py);
118        if self.is_text_io {
119            if buf.len() < 4 {
120                return Err(io::Error::new(
121                    io::ErrorKind::InvalidInput,
122                    "buffer size must be at least 4 bytes",
123                ));
124            }
125            let res = inner.call_method1(consts::read(py), (buf.len() / 4,))?;
126            let rust_string = res.extract::<Cow<str>>()?;
127            let bytes = rust_string.as_bytes();
128            buf.write_all(bytes)?;
129            Ok(bytes.len())
130        } else {
131            let pybytes = inner.call_method1(consts::read(py), (buf.len(),))?;
132            let bytes = pybytes.extract::<Cow<[u8]>>()?;
133            buf.write_all(&bytes)?;
134            Ok(bytes.len())
135        }
136    }
137
138    pub fn py_write(&self, py: Python<'_>, buf: &[u8]) -> io::Result<usize> {
139        let inner = self.inner.bind(py);
140        let arg = if self.is_text_io {
141            let s =
142                std::str::from_utf8(buf).expect("Tried to write non-utf8 data to a TextIO object.");
143            PyString::new(py, s).into_any()
144        } else {
145            PyBytes::new(py, buf).into_any()
146        };
147
148        let number_bytes_written = inner.call_method1(consts::write(py), (arg,))?;
149
150        if number_bytes_written.is_none() {
151            return Err(io::Error::new(
152                io::ErrorKind::Other,
153                "write() returned None, expected number of bytes written",
154            ));
155        }
156
157        number_bytes_written.extract().map_err(io::Error::from)
158    }
159
160    pub fn py_flush(&self, py: Python<'_>) -> io::Result<()> {
161        self.inner.call_method0(py, consts::flush(py))?;
162        Ok(())
163    }
164
165    pub fn py_seek(&self, py: Python<'_>, pos: SeekFrom) -> io::Result<u64> {
166        let inner = self.inner.bind(py);
167        let (whence, offset) = match pos {
168            SeekFrom::Start(offset) => (0, offset as i64),
169            SeekFrom::End(offset) => (2, offset),
170            SeekFrom::Current(offset) => (1, offset),
171        };
172
173        let res = inner.call_method1(consts::seek(py), (offset, whence))?;
174        res.extract().map_err(io::Error::from)
175    }
176
177    #[cfg(unix)]
178    pub fn py_as_raw_fd(&self, py: Python<'_>) -> RawFd {
179        let inner = self.inner.bind(py);
180        let fd = inner
181            .call_method0(consts::fileno(py))
182            .expect("Object does not have a fileno() method.");
183
184        fd.extract().expect("File descriptor is not an integer.")
185    }
186
187    pub fn py_clone(&self, py: Python<'_>) -> PyFileLikeObject {
188        PyFileLikeObject {
189            inner: self.inner.clone_ref(py),
190            is_text_io: self.is_text_io,
191        }
192    }
193
194    /// Access the name of the underlying file, if one exists
195    /// https://docs.python.org/3/library/io.html#io.FileIO.name
196    pub fn py_name(&self, py: Python<'_>) -> Option<String> {
197        let py_obj = self.inner.getattr(py, intern!(py, "name")).ok()?;
198        py_obj.extract::<String>(py).ok()
199    }
200}
201
202impl Read for PyFileLikeObject {
203    fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
204        Python::with_gil(|py| self.py_read(py, buf))
205    }
206}
207
208impl Read for &PyFileLikeObject {
209    fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
210        Python::with_gil(|py| self.py_read(py, buf))
211    }
212}
213
214impl Write for PyFileLikeObject {
215    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
216        Python::with_gil(|py| self.py_write(py, buf))
217    }
218
219    fn flush(&mut self) -> Result<(), io::Error> {
220        Python::with_gil(|py| self.py_flush(py))
221    }
222}
223
224impl Write for &PyFileLikeObject {
225    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
226        Python::with_gil(|py| self.py_write(py, buf))
227    }
228
229    fn flush(&mut self) -> Result<(), io::Error> {
230        Python::with_gil(|py| self.py_flush(py))
231    }
232}
233
234impl Seek for PyFileLikeObject {
235    fn seek(&mut self, pos: SeekFrom) -> Result<u64, io::Error> {
236        Python::with_gil(|py| self.py_seek(py, pos))
237    }
238}
239
240impl Seek for &PyFileLikeObject {
241    fn seek(&mut self, pos: SeekFrom) -> Result<u64, io::Error> {
242        Python::with_gil(|py| self.py_seek(py, pos))
243    }
244}
245
246#[cfg(unix)]
247impl AsRawFd for PyFileLikeObject {
248    fn as_raw_fd(&self) -> RawFd {
249        Python::with_gil(|py| self.py_as_raw_fd(py))
250    }
251}
252
253#[cfg(unix)]
254impl AsRawFd for &PyFileLikeObject {
255    fn as_raw_fd(&self) -> RawFd {
256        Python::with_gil(|py| self.py_as_raw_fd(py))
257    }
258}
259
260impl<'py> FromPyObject<'py> for PyFileLikeObject {
261    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
262        Self::py_new(ob.clone())
263    }
264}
265
266mod consts {
267    use pyo3::prelude::*;
268    use pyo3::sync::GILOnceCell;
269    use pyo3::types::PyString;
270    use pyo3::{intern, Bound, Py, PyResult, Python};
271
272    pub fn fileno(py: Python) -> &Bound<PyString> {
273        intern!(py, "fileno")
274    }
275
276    pub fn read(py: Python) -> &Bound<PyString> {
277        intern!(py, "read")
278    }
279
280    pub fn write(py: Python<'_>) -> &Bound<PyString> {
281        intern!(py, "write")
282    }
283
284    pub fn seek(py: Python<'_>) -> &Bound<PyString> {
285        intern!(py, "seek")
286    }
287
288    pub fn flush(py: Python<'_>) -> &Bound<PyString> {
289        intern!(py, "flush")
290    }
291
292    pub fn text_io_base(py: Python) -> PyResult<&Bound<PyAny>> {
293        static INSTANCE: GILOnceCell<Py<PyAny>> = GILOnceCell::new();
294
295        INSTANCE
296            .get_or_try_init(py, || {
297                let io = PyModule::import(py, "io")?;
298                let cls = io.getattr("TextIOBase")?;
299                Ok(cls.unbind())
300            })
301            .map(|x| x.bind(py))
302    }
303}
pyo3_file/lib.rs

pyo3_file/
lib.rs
