Implement PartialEq for PyBytes and [u8] (#4259)

* Copy pasta implementation from types/string.rs * changelog * I think I don't need a special implementation for 3.10 or ABI * Copy pasta tests * Fix comment with correct type Co-authored-by: Icxolu <[email protected]> * Fix implementation * Use slice in tests * Try renaming changelog file * Fix doc example * Fix doc example Co-authored-by: Icxolu <[email protected]> --------- Co-authored-by: Icxolu <[email protected]>
PyO3 · Jun 22, 2024 · 908ef6a · 908ef6a
1 parent c4d18e5
commit 908ef6a
Show file tree

Hide file tree

Showing 3 changed files with 160 additions and 0 deletions.
diff --git a/newsfragments/4250.added.md b/newsfragments/4250.added.md
@@ -0,0 +1 @@
+Implement `PartialEq<str>` for `Bound<'py, PyBytes>`.
diff --git a/newsfragments/4245.added.md → newsfragments/4259.added.md b/newsfragments/4245.added.md → newsfragments/4259.added.md
diff --git a/src/types/bytes.rs b/src/types/bytes.rs
@@ -11,6 +11,35 @@ use std::str;
 /// Represents a Python `bytes` object.
 ///
 /// This type is immutable.
+///
+/// # Equality
+///
+/// For convenience, [`Bound<'py, PyBytes>`] implements [`PartialEq<[u8]>`] to allow comparing the
+/// data in the Python bytes to a Rust `[u8]`.
+///
+/// This is not always the most appropriate way to compare Python bytes, as Python bytes subclasses
+/// may have different equality semantics. In situations where subclasses overriding equality might be
+/// relevant, use [`PyAnyMethods::eq`], at cost of the additional overhead of a Python method call.
+///
+/// ```rust
+/// # use pyo3::prelude::*;
+/// use pyo3::types::PyBytes;
+///
+/// # Python::with_gil(|py| {
+/// let py_bytes = PyBytes::new_bound(py, b"foo".as_slice());
+/// // via PartialEq<[u8]>
+/// assert_eq!(py_bytes, b"foo".as_slice());
+///
+/// // via Python equality
+/// let other = PyBytes::new_bound(py, b"foo".as_slice());
+/// assert!(py_bytes.as_any().eq(other).unwrap());
+///
+/// // Note that `eq` will convert it's argument to Python using `ToPyObject`,
+/// // so the following does not compare equal since the slice will convert into a
+/// // `list`, not a `bytes` object.
+/// assert!(!py_bytes.as_any().eq(b"foo".as_slice()).unwrap());
+/// # });
+/// ```
 #[repr(transparent)]
 pub struct PyBytes(PyAny);
 
@@ -191,6 +220,106 @@ impl<I: SliceIndex<[u8]>> Index<I> for Bound<'_, PyBytes> {
     }
 }
 
+/// Compares whether the Python bytes object is equal to the [u8].
+///
+/// In some cases Python equality might be more appropriate; see the note on [`PyBytes`].
+impl PartialEq<[u8]> for Bound<'_, PyBytes> {
+    #[inline]
+    fn eq(&self, other: &[u8]) -> bool {
+        self.as_borrowed() == *other
+    }
+}
+
+/// Compares whether the Python bytes object is equal to the [u8].
+///
+/// In some cases Python equality might be more appropriate; see the note on [`PyBytes`].
+impl PartialEq<&'_ [u8]> for Bound<'_, PyBytes> {
+    #[inline]
+    fn eq(&self, other: &&[u8]) -> bool {
+        self.as_borrowed() == **other
+    }
+}
+
+/// Compares whether the Python bytes object is equal to the [u8].
+///
+/// In some cases Python equality might be more appropriate; see the note on [`PyBytes`].
+impl PartialEq<Bound<'_, PyBytes>> for [u8] {
+    #[inline]
+    fn eq(&self, other: &Bound<'_, PyBytes>) -> bool {
+        *self == other.as_borrowed()
+    }
+}
+
+/// Compares whether the Python bytes object is equal to the [u8].
+///
+/// In some cases Python equality might be more appropriate; see the note on [`PyBytes`].
+impl PartialEq<&'_ Bound<'_, PyBytes>> for [u8] {
+    #[inline]
+    fn eq(&self, other: &&Bound<'_, PyBytes>) -> bool {
+        *self == other.as_borrowed()
+    }
+}
+
+/// Compares whether the Python bytes object is equal to the [u8].
+///
+/// In some cases Python equality might be more appropriate; see the note on [`PyBytes`].
+impl PartialEq<Bound<'_, PyBytes>> for &'_ [u8] {
+    #[inline]
+    fn eq(&self, other: &Bound<'_, PyBytes>) -> bool {
+        **self == other.as_borrowed()
+    }
+}
+
+/// Compares whether the Python bytes object is equal to the [u8].
+///
+/// In some cases Python equality might be more appropriate; see the note on [`PyBytes`].
+impl PartialEq<[u8]> for &'_ Bound<'_, PyBytes> {
+    #[inline]
+    fn eq(&self, other: &[u8]) -> bool {
+        self.as_borrowed() == other
+    }
+}
+
+/// Compares whether the Python bytes object is equal to the [u8].
+///
+/// In some cases Python equality might be more appropriate; see the note on [`PyBytes`].
+impl PartialEq<[u8]> for Borrowed<'_, '_, PyBytes> {
+    #[inline]
+    fn eq(&self, other: &[u8]) -> bool {
+        self.as_bytes() == other
+    }
+}
+
+/// Compares whether the Python bytes object is equal to the [u8].
+///
+/// In some cases Python equality might be more appropriate; see the note on [`PyBytes`].
+impl PartialEq<&[u8]> for Borrowed<'_, '_, PyBytes> {
+    #[inline]
+    fn eq(&self, other: &&[u8]) -> bool {
+        *self == **other
+    }
+}
+
+/// Compares whether the Python bytes object is equal to the [u8].
+///
+/// In some cases Python equality might be more appropriate; see the note on [`PyBytes`].
+impl PartialEq<Borrowed<'_, '_, PyBytes>> for [u8] {
+    #[inline]
+    fn eq(&self, other: &Borrowed<'_, '_, PyBytes>) -> bool {
+        other == self
+    }
+}
+
+/// Compares whether the Python bytes object is equal to the [u8].
+///
+/// In some cases Python equality might be more appropriate; see the note on [`PyBytes`].
+impl PartialEq<Borrowed<'_, '_, PyBytes>> for &'_ [u8] {
+    #[inline]
+    fn eq(&self, other: &Borrowed<'_, '_, PyBytes>) -> bool {
+        other == self
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -251,4 +380,34 @@ mod tests {
                 .is_instance_of::<PyValueError>(py));
         });
     }
+
+    #[test]
+    fn test_comparisons() {
+        Python::with_gil(|py| {
+            let b = b"hello, world".as_slice();
+            let py_bytes = PyBytes::new_bound(py, b);
+
+            assert_eq!(py_bytes, b"hello, world".as_slice());
+
+            assert_eq!(py_bytes, b);
+            assert_eq!(&py_bytes, b);
+            assert_eq!(b, py_bytes);
+            assert_eq!(b, &py_bytes);
+
+            assert_eq!(py_bytes, *b);
+            assert_eq!(&py_bytes, *b);
+            assert_eq!(*b, py_bytes);
+            assert_eq!(*b, &py_bytes);
+
+            let py_string = py_bytes.as_borrowed();
+
+            assert_eq!(py_string, b);
+            assert_eq!(&py_string, b);
+            assert_eq!(b, py_string);
+            assert_eq!(b, &py_string);
+
+            assert_eq!(py_string, *b);
+            assert_eq!(*b, py_string);
+        })
+    }
 }