pyo3/
pybacked.rs

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

pyo3/
pybacked.rs
