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

pyo3/
pybacked.rs
