pyo3/types/bytes/
writer.rs

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

pyo3/types/bytes/
writer.rs
