pyo3/
fmt.rs

1#[cfg(any(doc, all(Py_3_14, not(Py_LIMITED_API))))]
2use crate::{types::PyString, Python};
3#[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
4use {
5    crate::ffi::{
6        PyUnicodeWriter_Create, PyUnicodeWriter_Discard, PyUnicodeWriter_Finish,
7        PyUnicodeWriter_WriteChar, PyUnicodeWriter_WriteUTF8,
8    },
9    crate::ffi_ptr_ext::FfiPtrExt,
10    crate::impl_::callback::WrappingCastTo,
11    crate::py_result_ext::PyResultExt,
12    crate::IntoPyObject,
13    crate::{ffi, Bound, PyErr, PyResult},
14    std::fmt,
15    std::mem::ManuallyDrop,
16    std::ptr::NonNull,
17};
18
19/// This macro is analogous to Rust's [`format!`] macro, but returns a [`PyString`] instead of a [`String`].
20///
21/// # Arguments
22///
23/// The arguments are exactly like [`format!`], but with `py` (a [`Python`] token) as the first argument:
24///
25/// # Interning Advantage
26///
27/// If the format string is a static string and all arguments are constant at compile time,
28/// this macro will intern the string in Python, offering better performance and memory usage
29/// compared to [`PyString::from_fmt`].
30///
31/// ```rust
32/// # use pyo3::{py_format, Python, types::PyString, Bound};
33/// Python::attach(|py| {
34///     let py_string: Bound<'_, PyString> = py_format!(py, "{} {}", "hello", "world").unwrap();
35///     assert_eq!(py_string.to_string(), "hello world");
36/// });
37/// ```
38#[macro_export]
39macro_rules! py_format {
40    ($py: expr, $($arg:tt)*) => {{
41        if let Some(static_string) = format_args!($($arg)*).as_str() {
42            static INTERNED: $crate::sync::PyOnceLock<$crate::Py<$crate::types::PyString>> = $crate::sync::PyOnceLock::new();
43            Ok(
44                INTERNED
45                .get_or_init($py, || $crate::types::PyString::intern($py, static_string).unbind())
46                .bind($py)
47                .to_owned()
48            )
49        } else {
50            $crate::types::PyString::from_fmt($py, format_args!($($arg)*))
51        }
52    }}
53}
54
55#[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
56/// The `PyUnicodeWriter` is a utility for efficiently constructing Python strings
57pub(crate) struct PyUnicodeWriter<'py> {
58    python: Python<'py>,
59    writer: NonNull<ffi::PyUnicodeWriter>,
60    last_error: Option<PyErr>,
61}
62
63#[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
64impl<'py> PyUnicodeWriter<'py> {
65    /// Creates a new `PyUnicodeWriter`.
66    pub fn new(py: Python<'py>) -> PyResult<Self> {
67        Self::with_capacity(py, 0)
68    }
69
70    /// Creates a new `PyUnicodeWriter` with the specified initial capacity.
71    #[inline]
72    pub fn with_capacity(py: Python<'py>, capacity: usize) -> PyResult<Self> {
73        match NonNull::new(unsafe { PyUnicodeWriter_Create(capacity.wrapping_cast()) }) {
74            Some(ptr) => Ok(PyUnicodeWriter {
75                python: py,
76                writer: ptr,
77                last_error: None,
78            }),
79            None => Err(PyErr::fetch(py)),
80        }
81    }
82
83    /// Consumes the `PyUnicodeWriter` and returns a `Bound<PyString>` containing the constructed string.
84    #[inline]
85    pub fn into_py_string(mut self) -> PyResult<Bound<'py, PyString>> {
86        let py = self.python;
87        if let Some(error) = self.take_error() {
88            Err(error)
89        } else {
90            unsafe {
91                PyUnicodeWriter_Finish(ManuallyDrop::new(self).as_ptr())
92                    .assume_owned_or_err(py)
93                    .cast_into_unchecked()
94            }
95        }
96    }
97
98    /// When fmt::Write returned an error, this function can be used to retrieve the last error that occurred.
99    #[inline]
100    pub fn take_error(&mut self) -> Option<PyErr> {
101        self.last_error.take()
102    }
103
104    #[inline]
105    fn as_ptr(&self) -> *mut ffi::PyUnicodeWriter {
106        self.writer.as_ptr()
107    }
108
109    #[inline]
110    fn set_error(&mut self) {
111        self.last_error = Some(PyErr::fetch(self.python));
112    }
113}
114
115#[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
116impl fmt::Write for PyUnicodeWriter<'_> {
117    #[inline]
118    fn write_str(&mut self, s: &str) -> fmt::Result {
119        let result = unsafe {
120            PyUnicodeWriter_WriteUTF8(self.as_ptr(), s.as_ptr().cast(), s.len() as isize)
121        };
122        if result < 0 {
123            self.set_error();
124            Err(fmt::Error)
125        } else {
126            Ok(())
127        }
128    }
129
130    #[inline]
131    fn write_char(&mut self, c: char) -> fmt::Result {
132        let result = unsafe { PyUnicodeWriter_WriteChar(self.as_ptr(), c.into()) };
133        if result < 0 {
134            self.set_error();
135            Err(fmt::Error)
136        } else {
137            Ok(())
138        }
139    }
140}
141
142#[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
143impl Drop for PyUnicodeWriter<'_> {
144    #[inline]
145    fn drop(&mut self) {
146        unsafe {
147            PyUnicodeWriter_Discard(self.as_ptr());
148        }
149    }
150}
151
152#[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
153impl<'py> IntoPyObject<'py> for PyUnicodeWriter<'py> {
154    type Target = PyString;
155    type Output = Bound<'py, Self::Target>;
156    type Error = PyErr;
157
158    #[inline]
159    fn into_pyobject(self, _py: Python<'py>) -> PyResult<Bound<'py, PyString>> {
160        self.into_py_string()
161    }
162}
163
164#[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
165impl<'py> TryInto<Bound<'py, PyString>> for PyUnicodeWriter<'py> {
166    type Error = PyErr;
167
168    #[inline]
169    fn try_into(self) -> PyResult<Bound<'py, PyString>> {
170        self.into_py_string()
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    #[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
177    use super::*;
178    use crate::types::PyStringMethods;
179    use crate::{IntoPyObject, Python};
180
181    #[test]
182    #[allow(clippy::write_literal)]
183    #[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
184    fn unicode_writer_test() {
185        use std::fmt::Write;
186        Python::attach(|py| {
187            let mut writer = PyUnicodeWriter::new(py).unwrap();
188            write!(writer, "Hello {}!", "world").unwrap();
189            writer.write_char('😎').unwrap();
190            let result = writer.into_py_string().unwrap();
191            assert_eq!(result.to_string(), "Hello world!😎");
192        });
193    }
194
195    #[test]
196    #[allow(clippy::write_literal)]
197    #[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
198    fn unicode_writer_with_capacity() {
199        use std::fmt::Write;
200        Python::attach(|py| {
201            let mut writer = PyUnicodeWriter::with_capacity(py, 10).unwrap();
202            write!(writer, "Hello {}!", "world").unwrap();
203            writer.write_char('😎').unwrap();
204            let result = writer.into_py_string().unwrap();
205            assert_eq!(result.to_string(), "Hello world!😎");
206        });
207    }
208
209    #[test]
210    fn test_pystring_from_fmt() {
211        Python::attach(|py| {
212            py_format!(py, "Hello {}!", "world").unwrap();
213        });
214    }
215
216    #[test]
217    fn test_complex_format() {
218        Python::attach(|py| {
219            let complex_value = (42, "foo", [0; 0]).into_pyobject(py).unwrap();
220            let py_string = py_format!(py, "This is some complex value: {complex_value}").unwrap();
221            let actual = py_string.to_cow().unwrap();
222            let expected = "This is some complex value: (42, 'foo', [])";
223            assert_eq!(actual, expected);
224        });
225    }
226}
pyo3/fmt.rs

pyo3/
fmt.rs
