pyo3/
byteswriter.rs

1// TODO https://github.com/PyO3/pyo3/issues/5487
2#![allow(clippy::undocumented_unsafe_blocks)]
3
4#[cfg(feature = "experimental-inspect")]
5use crate::inspect::PyStaticExpr;
6#[allow(unused_imports, reason = "conditionally used")]
7use crate::platform::prelude::*;
8#[cfg(feature = "experimental-inspect")]
9use crate::PyTypeInfo;
10#[cfg(not(Py_LIMITED_API))]
11use crate::{
12    err::error_on_minusone,
13    ffi::{
14        self,
15        compat::{
16            PyBytesWriter_Create, PyBytesWriter_Discard, PyBytesWriter_Finish,
17            PyBytesWriter_GetData, PyBytesWriter_GetSize, PyBytesWriter_Resize,
18        },
19    },
20    ffi_ptr_ext::FfiPtrExt,
21    py_result_ext::PyResultExt,
22};
23use crate::{types::PyBytes, Bound, IntoPyObject, PyErr, PyResult, Python};
24#[cfg(not(Py_LIMITED_API))]
25use core::{
26    mem::ManuallyDrop,
27    ptr::{self, NonNull},
28};
29use std::io::IoSlice;
30
31pub struct PyBytesWriter<'py> {
32    python: Python<'py>,
33    #[cfg(not(Py_LIMITED_API))]
34    writer: NonNull<ffi::compat::PyBytesWriter>,
35    #[cfg(Py_LIMITED_API)]
36    buffer: Vec<u8>,
37}
38
39impl<'py> PyBytesWriter<'py> {
40    /// Create a new `PyBytesWriter` with a default initial capacity.
41    #[inline]
42    pub fn new(py: Python<'py>) -> PyResult<Self> {
43        Self::with_capacity(py, 0)
44    }
45
46    /// Create a new `PyBytesWriter` with the specified initial capacity.
47    #[inline]
48    #[cfg_attr(Py_LIMITED_API, allow(clippy::unnecessary_wraps))]
49    pub fn with_capacity(py: Python<'py>, capacity: usize) -> PyResult<Self> {
50        cfg_select! {
51            not(Py_LIMITED_API) => NonNull::new(unsafe { PyBytesWriter_Create(capacity as _) }).map_or_else(
52                || Err(PyErr::fetch(py)),
53                |writer| {
54                    let mut writer = PyBytesWriter { python: py, writer };
55                    if capacity > 0 {
56                        // SAFETY: By setting the length to 0, we ensure no bytes are considered uninitialized.
57                        unsafe { writer.set_len(0)? };
58                    }
59                    Ok(writer)
60                },
61            ),
62            Py_LIMITED_API => Ok(PyBytesWriter {
63                python: py,
64                buffer: Vec::with_capacity(capacity),
65            })
66        }
67    }
68
69    /// Get the current length of the internal buffer.
70    #[inline]
71    pub fn len(&self) -> usize {
72        cfg_select! {
73            not(Py_LIMITED_API) => unsafe {
74                PyBytesWriter_GetSize(self.writer.as_ptr()) as _
75            },
76            Py_LIMITED_API => self.buffer.len()
77        }
78    }
79
80    #[inline]
81    #[cfg(not(Py_LIMITED_API))]
82    fn as_mut_ptr(&mut self) -> *mut u8 {
83        unsafe { PyBytesWriter_GetData(self.writer.as_ptr()) as _ }
84    }
85
86    /// Set the length of the internal buffer to `new_len`. The new bytes are uninitialized.
87    ///
88    /// # Safety
89    /// The caller must ensure the new bytes are initialized. This will also make all pointers
90    /// returned by `as_mut_ptr` invalid, so the caller must not hold any references to the buffer
91    /// across this call.
92    #[inline]
93    #[cfg(not(Py_LIMITED_API))]
94    unsafe fn set_len(&mut self, new_len: usize) -> PyResult<()> {
95        unsafe {
96            error_on_minusone(
97                self.python,
98                PyBytesWriter_Resize(self.writer.as_ptr(), new_len as _),
99            )
100        }
101    }
102}
103
104impl<'py> TryFrom<PyBytesWriter<'py>> for Bound<'py, PyBytes> {
105    type Error = PyErr;
106
107    #[inline]
108    fn try_from(value: PyBytesWriter<'py>) -> Result<Self, Self::Error> {
109        let py = value.python;
110        cfg_select! {
111            // SAFETY:
112            //  - we no longer use `value` after this call
113            //  - `PyBytesWriter_Finish` will return a new reference to a bytes object on success
114            not(Py_LIMITED_API) => unsafe {
115                PyBytesWriter_Finish(ManuallyDrop::new(value).writer.as_ptr())
116                    .assume_owned_or_err(py)
117                    .cast_into_unchecked()
118            },
119            Py_LIMITED_API => Ok(PyBytes::new(py, &value.buffer))
120        }
121    }
122}
123
124impl<'py> IntoPyObject<'py> for PyBytesWriter<'py> {
125    type Target = PyBytes;
126    type Output = Bound<'py, PyBytes>;
127    type Error = PyErr;
128
129    #[cfg(feature = "experimental-inspect")]
130    const OUTPUT_TYPE: PyStaticExpr = PyBytes::TYPE_HINT;
131
132    #[inline]
133    fn into_pyobject(self, _py: Python<'py>) -> Result<Self::Output, Self::Error> {
134        self.try_into()
135    }
136}
137
138#[cfg(not(Py_LIMITED_API))]
139impl<'py> Drop for PyBytesWriter<'py> {
140    #[inline]
141    fn drop(&mut self) {
142        unsafe { PyBytesWriter_Discard(self.writer.as_ptr()) }
143    }
144}
145
146#[cfg(not(Py_LIMITED_API))]
147impl std::io::Write for PyBytesWriter<'_> {
148    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
149        self.write_all(buf)?;
150        Ok(buf.len())
151    }
152
153    fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> std::io::Result<usize> {
154        let len = bufs.iter().map(|b| b.len()).sum();
155        let pos = self.len();
156
157        // SAFETY: We write the new uninitialized bytes below.
158        unsafe { self.set_len(self.len() + len)? }
159
160        // SAFETY: We ensured enough capacity above and the ptr will be valid because we will not be
161        // resizing the buffer until we have written all the data.
162        let mut ptr = unsafe { self.as_mut_ptr().add(pos) };
163
164        for buf in bufs {
165            // SAFETY: We have ensured enough capacity above.
166            unsafe { ptr::copy_nonoverlapping(buf.as_ptr(), ptr, buf.len()) };
167
168            // SAFETY: We just wrote buf.len() bytes
169            ptr = unsafe { ptr.add(buf.len()) };
170        }
171        Ok(len)
172    }
173
174    fn flush(&mut self) -> std::io::Result<()> {
175        Ok(())
176    }
177
178    fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
179        let len = buf.len();
180        let pos = self.len();
181
182        // SAFETY: We write the new uninitialized bytes below.
183        unsafe { self.set_len(pos + len)? }
184
185        // SAFETY: We have ensured enough capacity above.
186        unsafe { ptr::copy_nonoverlapping(buf.as_ptr(), self.as_mut_ptr().add(pos), len) };
187
188        Ok(())
189    }
190}
191
192#[cfg(Py_LIMITED_API)]
193impl std::io::Write for PyBytesWriter<'_> {
194    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
195        self.buffer.write(buf)
196    }
197
198    fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> std::io::Result<usize> {
199        self.buffer.write_vectored(bufs)
200    }
201
202    fn flush(&mut self) -> std::io::Result<()> {
203        self.buffer.flush()
204    }
205
206    fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
207        self.buffer.write_all(buf)
208    }
209
210    fn write_fmt(&mut self, args: core::fmt::Arguments<'_>) -> std::io::Result<()> {
211        self.buffer.write_fmt(args)
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218    use crate::types::PyBytesMethods;
219    use std::io::Write;
220
221    #[test]
222    fn test_io_write() {
223        Python::attach(|py| {
224            let buf = b"hallo world";
225            let mut writer = PyBytesWriter::new(py).unwrap();
226            assert_eq!(writer.write(buf).unwrap(), 11);
227            let bytes: Bound<'_, PyBytes> = writer.try_into().unwrap();
228            assert_eq!(bytes.as_bytes(), buf);
229        })
230    }
231
232    #[test]
233    fn test_pre_allocated() {
234        Python::attach(|py| {
235            let buf = b"hallo world";
236            let mut writer = PyBytesWriter::with_capacity(py, buf.len()).unwrap();
237            assert_eq!(writer.len(), 0, "Writer position should be zero");
238            assert_eq!(writer.write(buf).unwrap(), 11);
239            let bytes: Bound<'_, PyBytes> = writer.try_into().unwrap();
240            assert_eq!(bytes.as_bytes(), buf);
241        })
242    }
243
244    #[test]
245    fn test_io_write_vectored() {
246        Python::attach(|py| {
247            let bufs = [IoSlice::new(b"hallo "), IoSlice::new(b"world")];
248            let mut writer = PyBytesWriter::new(py).unwrap();
249            assert_eq!(writer.write_vectored(&bufs).unwrap(), 11);
250            let bytes: Bound<'_, PyBytes> = writer.try_into().unwrap();
251            assert_eq!(bytes.as_bytes(), b"hallo world");
252        })
253    }
254
255    #[test]
256    fn test_io_write_vectored_large() {
257        Python::attach(|py| {
258            let large_data = vec![b'\n'; 1024]; // 1 KB
259            let bufs = [
260                IoSlice::new(b"hallo"),
261                IoSlice::new(&large_data),
262                IoSlice::new(b"world"),
263            ];
264            let mut writer = PyBytesWriter::new(py).unwrap();
265            assert_eq!(writer.write_vectored(&bufs).unwrap(), 1034);
266            let bytes: Bound<'_, PyBytes> = writer.try_into().unwrap();
267            assert!(bytes.as_bytes().starts_with(b"hallo\n"));
268            assert!(bytes.as_bytes().ends_with(b"world"));
269            assert_eq!(bytes.as_bytes().len(), 1034);
270        })
271    }
272
273    #[test]
274    fn test_large_data() {
275        Python::attach(|py| {
276            let mut writer = PyBytesWriter::new(py).unwrap();
277            let large_data = vec![0; 1024]; // 1 KB
278            writer.write_all(&large_data).unwrap();
279            let bytes: Bound<'_, PyBytes> = writer.try_into().unwrap();
280            assert_eq!(bytes.as_bytes(), large_data.as_slice());
281        })
282    }
283}
pyo3/byteswriter.rs

pyo3/
byteswriter.rs
