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::other(
152                "write() returned None, expected number of bytes written",
153            ));
154        }
155
156        number_bytes_written.extract().map_err(io::Error::from)
157    }
158
159    pub fn py_flush(&self, py: Python<'_>) -> io::Result<()> {
160        self.inner.call_method0(py, consts::flush(py))?;
161        Ok(())
162    }
163
164    pub fn py_seek(&self, py: Python<'_>, pos: SeekFrom) -> io::Result<u64> {
165        let inner = self.inner.bind(py);
166        let (whence, offset) = match pos {
167            SeekFrom::Start(offset) => (0, offset as i64),
168            SeekFrom::End(offset) => (2, offset),
169            SeekFrom::Current(offset) => (1, offset),
170        };
171
172        let res = inner.call_method1(consts::seek(py), (offset, whence))?;
173        res.extract().map_err(io::Error::from)
174    }
175
176    #[cfg(unix)]
177    pub fn py_as_raw_fd(&self, py: Python<'_>) -> RawFd {
178        let inner = self.inner.bind(py);
179        let fd = inner
180            .call_method0(consts::fileno(py))
181            .expect("Object does not have a fileno() method.");
182
183        fd.extract().expect("File descriptor is not an integer.")
184    }
185
186    pub fn py_clone(&self, py: Python<'_>) -> PyFileLikeObject {
187        PyFileLikeObject {
188            inner: self.inner.clone_ref(py),
189            is_text_io: self.is_text_io,
190        }
191    }
192
193    /// Access the name of the underlying file, if one exists
194    /// https://docs.python.org/3/library/io.html#io.FileIO.name
195    pub fn py_name(&self, py: Python<'_>) -> Option<String> {
196        let py_obj = self.inner.getattr(py, intern!(py, "name")).ok()?;
197        py_obj.extract::<String>(py).ok()
198    }
199}
200
201impl Read for PyFileLikeObject {
202    fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
203        Python::with_gil(|py| self.py_read(py, buf))
204    }
205}
206
207impl Read for &PyFileLikeObject {
208    fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
209        Python::with_gil(|py| self.py_read(py, buf))
210    }
211}
212
213impl Write for PyFileLikeObject {
214    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
215        Python::with_gil(|py| self.py_write(py, buf))
216    }
217
218    fn flush(&mut self) -> Result<(), io::Error> {
219        Python::with_gil(|py| self.py_flush(py))
220    }
221}
222
223impl Write for &PyFileLikeObject {
224    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
225        Python::with_gil(|py| self.py_write(py, buf))
226    }
227
228    fn flush(&mut self) -> Result<(), io::Error> {
229        Python::with_gil(|py| self.py_flush(py))
230    }
231}
232
233impl Seek for PyFileLikeObject {
234    fn seek(&mut self, pos: SeekFrom) -> Result<u64, io::Error> {
235        Python::with_gil(|py| self.py_seek(py, pos))
236    }
237}
238
239impl Seek for &PyFileLikeObject {
240    fn seek(&mut self, pos: SeekFrom) -> Result<u64, io::Error> {
241        Python::with_gil(|py| self.py_seek(py, pos))
242    }
243}
244
245#[cfg(unix)]
246impl AsRawFd for PyFileLikeObject {
247    fn as_raw_fd(&self) -> RawFd {
248        Python::with_gil(|py| self.py_as_raw_fd(py))
249    }
250}
251
252#[cfg(unix)]
253impl AsRawFd for &PyFileLikeObject {
254    fn as_raw_fd(&self) -> RawFd {
255        Python::with_gil(|py| self.py_as_raw_fd(py))
256    }
257}
258
259impl<'py> FromPyObject<'py> for PyFileLikeObject {
260    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
261        Self::py_new(ob.clone())
262    }
263}
264
265mod consts {
266    use pyo3::prelude::*;
267    use pyo3::sync::GILOnceCell;
268    use pyo3::types::PyString;
269    use pyo3::{intern, Bound, Py, PyResult, Python};
270
271    pub fn fileno(py: Python) -> &Bound<PyString> {
272        intern!(py, "fileno")
273    }
274
275    pub fn read(py: Python) -> &Bound<PyString> {
276        intern!(py, "read")
277    }
278
279    pub fn write(py: Python<'_>) -> &Bound<PyString> {
280        intern!(py, "write")
281    }
282
283    pub fn seek(py: Python<'_>) -> &Bound<PyString> {
284        intern!(py, "seek")
285    }
286
287    pub fn flush(py: Python<'_>) -> &Bound<PyString> {
288        intern!(py, "flush")
289    }
290
291    pub fn text_io_base(py: Python) -> PyResult<&Bound<PyAny>> {
292        static INSTANCE: GILOnceCell<Py<PyAny>> = GILOnceCell::new();
293
294        INSTANCE
295            .get_or_try_init(py, || {
296                let io = PyModule::import(py, "io")?;
297                let cls = io.getattr("TextIOBase")?;
298                Ok(cls.unbind())
299            })
300            .map(|x| x.bind(py))
301    }
302}
pyo3_file/lib.rs

pyo3_file/
lib.rs
