pyo3/
pybacked.rs

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

pyo3/
pybacked.rs
