Provide a special purpose FromPyObject impl to efficiently and safely get a byte slice from either bytes or byte arrays.

newsfragments/2899.added.md

@@ -0,0 +1 @@
+Add special-purpose impl of `FromPyObject` for `Cow<[u8]>` to efficiently handle both `bytes` and `bytearray` objects

src/types/bytes.rs

@@ -1,9 +1,12 @@
-use crate::{ffi, AsPyPointer, Py, PyAny, PyResult, Python};
+use crate::{ffi, AsPyPointer, FromPyObject, Py, PyAny, PyResult, Python};
+use std::borrow::Cow;
 use std::ops::Index;
 use std::os::raw::c_char;
 use std::slice::SliceIndex;
 use std::str;
 
+use super::bytearray::PyByteArray;
+
 /// Represents a Python `bytes` object.
 ///
 /// This type is immutable.
@@ -121,11 +124,25 @@ impl<I: SliceIndex<[u8]>> Index<I> for PyBytes {
     }
 }
 
+/// Special-purpose trait impl to efficiently handle both `bytes` and `bytearray`
+///
+/// If the source object is a `bytes` object, the `Cow` will be borrowed and
+/// pointing into the source object, and no copying or heap allocations will happen.
+/// If it is a `bytearray`, its contents will be copied to an owned `Cow`.
+impl<'source> FromPyObject<'source> for Cow<'source, [u8]> {
+    fn extract(ob: &'source PyAny) -> PyResult<Self> {
+        if let Ok(bytes) = ob.downcast::<PyBytes>() {
+            return Ok(Cow::Borrowed(bytes.as_bytes()));
+        }
+
+        let byte_array = ob.downcast::<PyByteArray>()?;
+        Ok(Cow::Owned(byte_array.to_vec()))
+    }
+}
+
 #[cfg(test)]
 mod tests {
-    use super::PyBytes;
-    use crate::FromPyObject;
-    use crate::Python;
+    use super::*;
 
     #[test]
     fn test_bytes_index() {
@@ -172,4 +189,22 @@ mod tests {
                 .is_instance_of::<PyValueError>(py));
         });
     }
+
+    #[test]
+    fn test_cow_impl() {
+        Python::with_gil(|py| {
+            let bytes = py.eval(r#"b"foobar""#, None, None).unwrap();
+            let cow = bytes.extract::<Cow<'_, [u8]>>().unwrap();
+            assert_eq!(cow, Cow::<[u8]>::Borrowed(b"foobar"));
+
+            let byte_array = py.eval(r#"bytearray(b"foobar")"#, None, None).unwrap();
+            let cow = byte_array.extract::<Cow<'_, [u8]>>().unwrap();
+            assert_eq!(cow, Cow::<[u8]>::Owned(b"foobar".to_vec()));
+
+            let something_else_entirely = py.eval("42", None, None).unwrap();
+            something_else_entirely
+                .extract::<Cow<'_, [u8]>>()
+                .unwrap_err();
+        });
+    }
 }