diff --git a/Cargo.toml b/Cargo.toml index b704ae368a2..1f90efd801a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,6 +34,7 @@ inventory = { version = "0.3.5", optional = true } # crate integrations that can be added using the eponymous features anyhow = { version = "1.0.1", optional = true } bigdecimal = {version = "0.4", optional = true } +bytes = { version = "1.10.1", optional = true } chrono = { version = "0.4.25", default-features = false, optional = true } chrono-tz = { version = ">= 0.10, < 0.11", default-features = false, optional = true } either = { version = "1.9", optional = true } @@ -125,6 +126,7 @@ full = [ # "multiple-pymethods", # Not supported by wasm "anyhow", "bigdecimal", + "bytes", "chrono", "chrono-tz", "either", diff --git a/guide/src/features.md b/guide/src/features.md index 971a60b0e25..9b9114157d4 100644 --- a/guide/src/features.md +++ b/guide/src/features.md @@ -124,6 +124,9 @@ Adds a dependency on [anyhow](https://docs.rs/anyhow). Enables a conversion from ### `bigdecimal` Adds a dependency on [bigdecimal](https://docs.rs/bigdecimal) and enables conversions into its [`BigDecimal`](https://docs.rs/bigdecimal/latest/bigdecimal/struct.BigDecimal.html) type. +### `bytes` +Adds a dependency on [bytes](https://docs.rs/bytes/latest/bytes) and enables conversions into its [`Bytes`](https://docs.rs/bytes/latest/bytes/struct.Bytes.html) and [`BytesMut`](https://docs.rs/bytes/latest/bytes/struct.BytesMut.html) types. + ### `chrono` Adds a dependency on [chrono](https://docs.rs/chrono). Enables a conversion from [chrono](https://docs.rs/chrono)'s types to python: diff --git a/newsfragments/5111.added.md b/newsfragments/5111.added.md new file mode 100644 index 00000000000..3d9362e6c05 --- /dev/null +++ b/newsfragments/5111.added.md @@ -0,0 +1 @@ +Add bytes to/from python conversions. diff --git a/src/conversions/bytes.rs b/src/conversions/bytes.rs new file mode 100644 index 00000000000..2ee79ef0ab0 --- /dev/null +++ b/src/conversions/bytes.rs @@ -0,0 +1,161 @@ +#![cfg(feature = "bytes")] + +//! Conversions to and from [bytes](https://docs.rs/bytes/latest/bytes/)'s [`Bytes`] and +//! [`BytesMut`] types. +//! +//! This is useful for efficiently converting Python's `bytes` and `bytearray` types efficiently. +//! +//! # Setup +//! +//! To use this feature, add in your **`Cargo.toml`**: +//! +//! ```toml +//! [dependencies] +#![doc = concat!("pyo3 = { version = \"", env!("CARGO_PKG_VERSION"), "\", features = [\"bytes\"] }")] +//! bytes = "1.10" +//! +//! Note that you must use compatible versions of bytes and PyO3. +//! +//! # Example +//! +//! Rust code to create functions which return `Bytes` or take `Bytes` as arguements: +//! +//! ```rust,no_run +//! use pyo3::prelude::*; +//! +//! #[pyfunction] +//! fn get_message_bytes() -> Bytes { +//! Bytes::from(b"Hello Python!".to_vec()) +//! } +//! +//! #[pyfunction] +//! fn num_bytes(bytes: Bytes) -> usize { +//! bytes.len() +//! } +//! +//! #[pymodule] +//! fn my_module(m: &Bound<'_, PyModule>) -> PyResult<()> { +//! m.add_function(wrap_pyfunction!(get_message_bytes, m)?)?; +//! m.add_function(wrap_pyfunction!(num_bytes, m)?)?; +//! Ok(()) +//! } +//! ``` +//! +//! Python code that calls these functions: +//! +//! ```python +//! from my_module import get_message_bytes, num_bytes +//! +//! message = get_message_bytes() +//! assert message == b"Hello Python!" +//! +//! size = num_bytes(message) +//! assert size == 13 +//! ``` +use bytes::{Bytes, BytesMut}; + +use crate::conversion::IntoPyObject; +use crate::exceptions::PyTypeError; +use crate::instance::Bound; +use crate::types::any::PyAnyMethods; +use crate::types::{PyByteArray, PyByteArrayMethods, PyBytes, PyBytesMethods}; +use crate::{FromPyObject, PyAny, PyErr, PyResult, Python}; + +impl FromPyObject<'_> for Bytes { + fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult { + if let Ok(bytes) = ob.downcast::() { + Ok(Bytes::from((*bytes).as_bytes().to_vec())) + } else if let Ok(bytearray) = ob.downcast::() { + Ok(Bytes::from((*bytearray).to_vec())) + } else { + Err(PyTypeError::new_err("expected bytes or bytearray")) + } + } +} + +impl FromPyObject<'_> for BytesMut { + fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult { + let bytes = ob.extract::()?; + Ok(BytesMut::from(bytes)) + } +} + +impl<'py> IntoPyObject<'py> for Bytes { + type Target = PyBytes; + type Output = Bound<'py, Self::Target>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(PyBytes::new(py, &self)) + } +} + +impl<'py> IntoPyObject<'py> for BytesMut { + type Target = PyBytes; + type Output = Bound<'py, Self::Target>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(PyBytes::new(py, &self)) + } +} + +#[cfg(test)] +mod tests { + use bytes::{Bytes, BytesMut}; + + use crate::{ + conversion::IntoPyObject, + ffi, + types::{any::PyAnyMethods, PyBytes}, + Python, + }; + + #[test] + fn test_bytes() { + Python::with_gil(|py| { + let py_bytes = py.eval(ffi::c_str!("b'foobar'"), None, None).unwrap(); + let bytes: Bytes = py_bytes.extract().unwrap(); + assert_eq!(bytes, Bytes::from(b"foobar".to_vec())); + + let bytes = Bytes::from(b"foobar".to_vec()).into_pyobject(py).unwrap(); + assert!(bytes.is_instance_of::()); + }); + } + + #[test] + fn test_bytearray() { + Python::with_gil(|py| { + let py_bytearray = py + .eval(ffi::c_str!("bytearray(b'foobar')"), None, None) + .unwrap(); + let bytes: Bytes = py_bytearray.extract().unwrap(); + assert_eq!(bytes, Bytes::from(b"foobar".to_vec())); + }); + } + + #[test] + fn test_bytes_mut() { + Python::with_gil(|py| { + let py_bytearray = py + .eval(ffi::c_str!("bytearray(b'foobar')"), None, None) + .unwrap(); + let bytes: BytesMut = py_bytearray.extract().unwrap(); + assert_eq!(bytes, BytesMut::from(&b"foobar"[..])); + + let bytesmut = BytesMut::from(&b"foobar"[..]).into_pyobject(py).unwrap(); + assert!(bytesmut.is_instance_of::()); + }); + } + + #[test] + fn test_bytearray_mut() { + Python::with_gil(|py| { + let py_bytearray = py + .eval(ffi::c_str!("bytearray(b'foobar')"), None, None) + .unwrap(); + let bytes: BytesMut = py_bytearray.extract().unwrap(); + assert_eq!(bytes, BytesMut::from(&b"foobar"[..])); + }); + } +} diff --git a/src/conversions/mod.rs b/src/conversions/mod.rs index 5403916ba8d..69ba77f6062 100644 --- a/src/conversions/mod.rs +++ b/src/conversions/mod.rs @@ -2,6 +2,7 @@ pub mod anyhow; pub mod bigdecimal; +pub mod bytes; pub mod chrono; pub mod chrono_tz; pub mod either;