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

pyo3_file/
lib.rs
