pyo3/
pybacked.rs

1//! Contains types for working with Python objects that own the underlying data.
2
3use std::{ops::Deref, ptr::NonNull, sync::Arc};
4
5use crate::{
6    types::{
7        any::PyAnyMethods, bytearray::PyByteArrayMethods, bytes::PyBytesMethods,
8        string::PyStringMethods, PyByteArray, PyBytes, PyString,
9    },
10    Bound, DowncastError, FromPyObject, IntoPy, Py, PyAny, PyErr, PyResult, Python, ToPyObject,
11};
12
13/// A wrapper around `str` where the storage is owned by a Python `bytes` or `str` object.
14///
15/// This type gives access to the underlying data via a `Deref` implementation.
16#[cfg_attr(feature = "py-clone", derive(Clone))]
17pub struct PyBackedStr {
18    #[allow(dead_code)] // only held so that the storage is not dropped
19    storage: Py<PyAny>,
20    data: NonNull<str>,
21}
22
23impl Deref for PyBackedStr {
24    type Target = str;
25    fn deref(&self) -> &str {
26        // Safety: `data` is known to be immutable and owned by self
27        unsafe { self.data.as_ref() }
28    }
29}
30
31impl AsRef<str> for PyBackedStr {
32    fn as_ref(&self) -> &str {
33        self
34    }
35}
36
37impl AsRef<[u8]> for PyBackedStr {
38    fn as_ref(&self) -> &[u8] {
39        self.as_bytes()
40    }
41}
42
43// Safety: the underlying Python str (or bytes) is immutable and
44// safe to share between threads
45unsafe impl Send for PyBackedStr {}
46unsafe impl Sync for PyBackedStr {}
47
48impl std::fmt::Display for PyBackedStr {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        self.deref().fmt(f)
51    }
52}
53
54impl_traits!(PyBackedStr, str);
55
56impl TryFrom<Bound<'_, PyString>> for PyBackedStr {
57    type Error = PyErr;
58    fn try_from(py_string: Bound<'_, PyString>) -> Result<Self, Self::Error> {
59        #[cfg(any(Py_3_10, not(Py_LIMITED_API)))]
60        {
61            let s = py_string.to_str()?;
62            let data = NonNull::from(s);
63            Ok(Self {
64                storage: py_string.as_any().to_owned().unbind(),
65                data,
66            })
67        }
68        #[cfg(not(any(Py_3_10, not(Py_LIMITED_API))))]
69        {
70            let bytes = py_string.encode_utf8()?;
71            let s = unsafe { std::str::from_utf8_unchecked(bytes.as_bytes()) };
72            let data = NonNull::from(s);
73            Ok(Self {
74                storage: bytes.into_any().unbind(),
75                data,
76            })
77        }
78    }
79}
80
81impl FromPyObject<'_> for PyBackedStr {
82    fn extract_bound(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
83        let py_string = obj.downcast::<PyString>()?.to_owned();
84        Self::try_from(py_string)
85    }
86}
87
88impl ToPyObject for PyBackedStr {
89    #[cfg(any(Py_3_10, not(Py_LIMITED_API)))]
90    fn to_object(&self, py: Python<'_>) -> Py<PyAny> {
91        self.storage.clone_ref(py)
92    }
93    #[cfg(not(any(Py_3_10, not(Py_LIMITED_API))))]
94    fn to_object(&self, py: Python<'_>) -> Py<PyAny> {
95        PyString::new_bound(py, self).into_any().unbind()
96    }
97}
98
99impl IntoPy<Py<PyAny>> for PyBackedStr {
100    #[cfg(any(Py_3_10, not(Py_LIMITED_API)))]
101    fn into_py(self, _py: Python<'_>) -> Py<PyAny> {
102        self.storage
103    }
104    #[cfg(not(any(Py_3_10, not(Py_LIMITED_API))))]
105    fn into_py(self, py: Python<'_>) -> Py<PyAny> {
106        PyString::new_bound(py, &self).into_any().unbind()
107    }
108}
109
110/// A wrapper around `[u8]` where the storage is either owned by a Python `bytes` object, or a Rust `Box<[u8]>`.
111///
112/// This type gives access to the underlying data via a `Deref` implementation.
113#[cfg_attr(feature = "py-clone", derive(Clone))]
114pub struct PyBackedBytes {
115    #[allow(dead_code)] // only held so that the storage is not dropped
116    storage: PyBackedBytesStorage,
117    data: NonNull<[u8]>,
118}
119
120#[allow(dead_code)]
121#[cfg_attr(feature = "py-clone", derive(Clone))]
122enum PyBackedBytesStorage {
123    Python(Py<PyBytes>),
124    Rust(Arc<[u8]>),
125}
126
127impl Deref for PyBackedBytes {
128    type Target = [u8];
129    fn deref(&self) -> &[u8] {
130        // Safety: `data` is known to be immutable and owned by self
131        unsafe { self.data.as_ref() }
132    }
133}
134
135impl AsRef<[u8]> for PyBackedBytes {
136    fn as_ref(&self) -> &[u8] {
137        self
138    }
139}
140
141// Safety: the underlying Python bytes or Rust bytes is immutable and
142// safe to share between threads
143unsafe impl Send for PyBackedBytes {}
144unsafe impl Sync for PyBackedBytes {}
145
146impl<const N: usize> PartialEq<[u8; N]> for PyBackedBytes {
147    fn eq(&self, other: &[u8; N]) -> bool {
148        self.deref() == other
149    }
150}
151
152impl<const N: usize> PartialEq<PyBackedBytes> for [u8; N] {
153    fn eq(&self, other: &PyBackedBytes) -> bool {
154        self == other.deref()
155    }
156}
157
158impl<const N: usize> PartialEq<&[u8; N]> for PyBackedBytes {
159    fn eq(&self, other: &&[u8; N]) -> bool {
160        self.deref() == *other
161    }
162}
163
164impl<const N: usize> PartialEq<PyBackedBytes> for &[u8; N] {
165    fn eq(&self, other: &PyBackedBytes) -> bool {
166        self == &other.deref()
167    }
168}
169
170impl_traits!(PyBackedBytes, [u8]);
171
172impl From<Bound<'_, PyBytes>> for PyBackedBytes {
173    fn from(py_bytes: Bound<'_, PyBytes>) -> Self {
174        let b = py_bytes.as_bytes();
175        let data = NonNull::from(b);
176        Self {
177            storage: PyBackedBytesStorage::Python(py_bytes.to_owned().unbind()),
178            data,
179        }
180    }
181}
182
183impl From<Bound<'_, PyByteArray>> for PyBackedBytes {
184    fn from(py_bytearray: Bound<'_, PyByteArray>) -> Self {
185        let s = Arc::<[u8]>::from(py_bytearray.to_vec());
186        let data = NonNull::from(s.as_ref());
187        Self {
188            storage: PyBackedBytesStorage::Rust(s),
189            data,
190        }
191    }
192}
193
194impl FromPyObject<'_> for PyBackedBytes {
195    fn extract_bound(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
196        if let Ok(bytes) = obj.downcast::<PyBytes>() {
197            Ok(Self::from(bytes.to_owned()))
198        } else if let Ok(bytearray) = obj.downcast::<PyByteArray>() {
199            Ok(Self::from(bytearray.to_owned()))
200        } else {
201            Err(DowncastError::new(obj, "`bytes` or `bytearray`").into())
202        }
203    }
204}
205
206impl ToPyObject for PyBackedBytes {
207    fn to_object(&self, py: Python<'_>) -> Py<PyAny> {
208        match &self.storage {
209            PyBackedBytesStorage::Python(bytes) => bytes.to_object(py),
210            PyBackedBytesStorage::Rust(bytes) => PyBytes::new_bound(py, bytes).into_any().unbind(),
211        }
212    }
213}
214
215impl IntoPy<Py<PyAny>> for PyBackedBytes {
216    fn into_py(self, py: Python<'_>) -> Py<PyAny> {
217        match self.storage {
218            PyBackedBytesStorage::Python(bytes) => bytes.into_any(),
219            PyBackedBytesStorage::Rust(bytes) => PyBytes::new_bound(py, &bytes).into_any().unbind(),
220        }
221    }
222}
223
224macro_rules! impl_traits {
225    ($slf:ty, $equiv:ty) => {
226        impl std::fmt::Debug for $slf {
227            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
228                self.deref().fmt(f)
229            }
230        }
231
232        impl PartialEq for $slf {
233            fn eq(&self, other: &Self) -> bool {
234                self.deref() == other.deref()
235            }
236        }
237
238        impl PartialEq<$equiv> for $slf {
239            fn eq(&self, other: &$equiv) -> bool {
240                self.deref() == other
241            }
242        }
243
244        impl PartialEq<&$equiv> for $slf {
245            fn eq(&self, other: &&$equiv) -> bool {
246                self.deref() == *other
247            }
248        }
249
250        impl PartialEq<$slf> for $equiv {
251            fn eq(&self, other: &$slf) -> bool {
252                self == other.deref()
253            }
254        }
255
256        impl PartialEq<$slf> for &$equiv {
257            fn eq(&self, other: &$slf) -> bool {
258                self == &other.deref()
259            }
260        }
261
262        impl Eq for $slf {}
263
264        impl PartialOrd for $slf {
265            fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
266                Some(self.cmp(other))
267            }
268        }
269
270        impl PartialOrd<$equiv> for $slf {
271            fn partial_cmp(&self, other: &$equiv) -> Option<std::cmp::Ordering> {
272                self.deref().partial_cmp(other)
273            }
274        }
275
276        impl PartialOrd<$slf> for $equiv {
277            fn partial_cmp(&self, other: &$slf) -> Option<std::cmp::Ordering> {
278                self.partial_cmp(other.deref())
279            }
280        }
281
282        impl Ord for $slf {
283            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
284                self.deref().cmp(other.deref())
285            }
286        }
287
288        impl std::hash::Hash for $slf {
289            fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
290                self.deref().hash(state)
291            }
292        }
293    };
294}
295use impl_traits;
296
297#[cfg(test)]
298mod test {
299    use super::*;
300    use crate::Python;
301    use std::collections::hash_map::DefaultHasher;
302    use std::hash::{Hash, Hasher};
303
304    #[test]
305    fn py_backed_str_empty() {
306        Python::with_gil(|py| {
307            let s = PyString::new_bound(py, "");
308            let py_backed_str = s.extract::<PyBackedStr>().unwrap();
309            assert_eq!(&*py_backed_str, "");
310        });
311    }
312
313    #[test]
314    fn py_backed_str() {
315        Python::with_gil(|py| {
316            let s = PyString::new_bound(py, "hello");
317            let py_backed_str = s.extract::<PyBackedStr>().unwrap();
318            assert_eq!(&*py_backed_str, "hello");
319        });
320    }
321
322    #[test]
323    fn py_backed_str_try_from() {
324        Python::with_gil(|py| {
325            let s = PyString::new_bound(py, "hello");
326            let py_backed_str = PyBackedStr::try_from(s).unwrap();
327            assert_eq!(&*py_backed_str, "hello");
328        });
329    }
330
331    #[test]
332    fn py_backed_str_to_object() {
333        Python::with_gil(|py| {
334            let orig_str = PyString::new_bound(py, "hello");
335            let py_backed_str = orig_str.extract::<PyBackedStr>().unwrap();
336            let new_str = py_backed_str.to_object(py);
337            assert_eq!(new_str.extract::<PyBackedStr>(py).unwrap(), "hello");
338            #[cfg(any(Py_3_10, not(Py_LIMITED_API)))]
339            assert!(new_str.is(&orig_str));
340        });
341    }
342
343    #[test]
344    fn py_backed_str_into_py() {
345        Python::with_gil(|py| {
346            let orig_str = PyString::new_bound(py, "hello");
347            let py_backed_str = orig_str.extract::<PyBackedStr>().unwrap();
348            let new_str = py_backed_str.into_py(py);
349            assert_eq!(new_str.extract::<PyBackedStr>(py).unwrap(), "hello");
350            #[cfg(any(Py_3_10, not(Py_LIMITED_API)))]
351            assert!(new_str.is(&orig_str));
352        });
353    }
354
355    #[test]
356    fn py_backed_bytes_empty() {
357        Python::with_gil(|py| {
358            let b = PyBytes::new_bound(py, b"");
359            let py_backed_bytes = b.extract::<PyBackedBytes>().unwrap();
360            assert_eq!(&*py_backed_bytes, b"");
361        });
362    }
363
364    #[test]
365    fn py_backed_bytes() {
366        Python::with_gil(|py| {
367            let b = PyBytes::new_bound(py, b"abcde");
368            let py_backed_bytes = b.extract::<PyBackedBytes>().unwrap();
369            assert_eq!(&*py_backed_bytes, b"abcde");
370        });
371    }
372
373    #[test]
374    fn py_backed_bytes_from_bytes() {
375        Python::with_gil(|py| {
376            let b = PyBytes::new_bound(py, b"abcde");
377            let py_backed_bytes = PyBackedBytes::from(b);
378            assert_eq!(&*py_backed_bytes, b"abcde");
379        });
380    }
381
382    #[test]
383    fn py_backed_bytes_from_bytearray() {
384        Python::with_gil(|py| {
385            let b = PyByteArray::new_bound(py, b"abcde");
386            let py_backed_bytes = PyBackedBytes::from(b);
387            assert_eq!(&*py_backed_bytes, b"abcde");
388        });
389    }
390
391    #[test]
392    fn py_backed_bytes_into_py() {
393        Python::with_gil(|py| {
394            let orig_bytes = PyBytes::new_bound(py, b"abcde");
395            let py_backed_bytes = PyBackedBytes::from(orig_bytes.clone());
396            assert!(py_backed_bytes.to_object(py).is(&orig_bytes));
397            assert!(py_backed_bytes.into_py(py).is(&orig_bytes));
398        });
399    }
400
401    #[test]
402    fn rust_backed_bytes_into_py() {
403        Python::with_gil(|py| {
404            let orig_bytes = PyByteArray::new_bound(py, b"abcde");
405            let rust_backed_bytes = PyBackedBytes::from(orig_bytes);
406            assert!(matches!(
407                rust_backed_bytes.storage,
408                PyBackedBytesStorage::Rust(_)
409            ));
410            let to_object = rust_backed_bytes.to_object(py).into_bound(py);
411            assert!(&to_object.is_exact_instance_of::<PyBytes>());
412            assert_eq!(&to_object.extract::<PyBackedBytes>().unwrap(), b"abcde");
413            let into_py = rust_backed_bytes.into_py(py).into_bound(py);
414            assert!(&into_py.is_exact_instance_of::<PyBytes>());
415            assert_eq!(&into_py.extract::<PyBackedBytes>().unwrap(), b"abcde");
416        });
417    }
418
419    #[test]
420    fn test_backed_types_send_sync() {
421        fn is_send<T: Send>() {}
422        fn is_sync<T: Sync>() {}
423
424        is_send::<PyBackedStr>();
425        is_sync::<PyBackedStr>();
426
427        is_send::<PyBackedBytes>();
428        is_sync::<PyBackedBytes>();
429    }
430
431    #[cfg(feature = "py-clone")]
432    #[test]
433    fn test_backed_str_clone() {
434        Python::with_gil(|py| {
435            let s1: PyBackedStr = PyString::new_bound(py, "hello").try_into().unwrap();
436            let s2 = s1.clone();
437            assert_eq!(s1, s2);
438
439            drop(s1);
440            assert_eq!(s2, "hello");
441        });
442    }
443
444    #[test]
445    fn test_backed_str_eq() {
446        Python::with_gil(|py| {
447            let s1: PyBackedStr = PyString::new_bound(py, "hello").try_into().unwrap();
448            let s2: PyBackedStr = PyString::new_bound(py, "hello").try_into().unwrap();
449            assert_eq!(s1, "hello");
450            assert_eq!(s1, s2);
451
452            let s3: PyBackedStr = PyString::new_bound(py, "abcde").try_into().unwrap();
453            assert_eq!("abcde", s3);
454            assert_ne!(s1, s3);
455        });
456    }
457
458    #[test]
459    fn test_backed_str_hash() {
460        Python::with_gil(|py| {
461            let h = {
462                let mut hasher = DefaultHasher::new();
463                "abcde".hash(&mut hasher);
464                hasher.finish()
465            };
466
467            let s1: PyBackedStr = PyString::new_bound(py, "abcde").try_into().unwrap();
468            let h1 = {
469                let mut hasher = DefaultHasher::new();
470                s1.hash(&mut hasher);
471                hasher.finish()
472            };
473
474            assert_eq!(h, h1);
475        });
476    }
477
478    #[test]
479    fn test_backed_str_ord() {
480        Python::with_gil(|py| {
481            let mut a = vec!["a", "c", "d", "b", "f", "g", "e"];
482            let mut b = a
483                .iter()
484                .map(|s| PyString::new_bound(py, s).try_into().unwrap())
485                .collect::<Vec<PyBackedStr>>();
486
487            a.sort();
488            b.sort();
489
490            assert_eq!(a, b);
491        })
492    }
493
494    #[cfg(feature = "py-clone")]
495    #[test]
496    fn test_backed_bytes_from_bytes_clone() {
497        Python::with_gil(|py| {
498            let b1: PyBackedBytes = PyBytes::new_bound(py, b"abcde").into();
499            let b2 = b1.clone();
500            assert_eq!(b1, b2);
501
502            drop(b1);
503            assert_eq!(b2, b"abcde");
504        });
505    }
506
507    #[cfg(feature = "py-clone")]
508    #[test]
509    fn test_backed_bytes_from_bytearray_clone() {
510        Python::with_gil(|py| {
511            let b1: PyBackedBytes = PyByteArray::new_bound(py, b"abcde").into();
512            let b2 = b1.clone();
513            assert_eq!(b1, b2);
514
515            drop(b1);
516            assert_eq!(b2, b"abcde");
517        });
518    }
519
520    #[test]
521    fn test_backed_bytes_eq() {
522        Python::with_gil(|py| {
523            let b1: PyBackedBytes = PyBytes::new_bound(py, b"abcde").into();
524            let b2: PyBackedBytes = PyByteArray::new_bound(py, b"abcde").into();
525
526            assert_eq!(b1, b"abcde");
527            assert_eq!(b1, b2);
528
529            let b3: PyBackedBytes = PyBytes::new_bound(py, b"hello").into();
530            assert_eq!(b"hello", b3);
531            assert_ne!(b1, b3);
532        });
533    }
534
535    #[test]
536    fn test_backed_bytes_hash() {
537        Python::with_gil(|py| {
538            let h = {
539                let mut hasher = DefaultHasher::new();
540                b"abcde".hash(&mut hasher);
541                hasher.finish()
542            };
543
544            let b1: PyBackedBytes = PyBytes::new_bound(py, b"abcde").into();
545            let h1 = {
546                let mut hasher = DefaultHasher::new();
547                b1.hash(&mut hasher);
548                hasher.finish()
549            };
550
551            let b2: PyBackedBytes = PyByteArray::new_bound(py, b"abcde").into();
552            let h2 = {
553                let mut hasher = DefaultHasher::new();
554                b2.hash(&mut hasher);
555                hasher.finish()
556            };
557
558            assert_eq!(h, h1);
559            assert_eq!(h, h2);
560        });
561    }
562
563    #[test]
564    fn test_backed_bytes_ord() {
565        Python::with_gil(|py| {
566            let mut a = vec![b"a", b"c", b"d", b"b", b"f", b"g", b"e"];
567            let mut b = a
568                .iter()
569                .map(|&b| PyBytes::new_bound(py, b).into())
570                .collect::<Vec<PyBackedBytes>>();
571
572            a.sort();
573            b.sort();
574
575            assert_eq!(a, b);
576        })
577    }
578}
pyo3/pybacked.rs

pyo3/
pybacked.rs
