string: implement API to access raw string data

With the recent implementation of non-limited unicode APIs, we're
able to query Python's low-level state to access the raw bytes that
Python is using to store string objects.

This commit implements a safe Rust API for obtaining a view into
Python's internals and representing the raw bytes Python is using
to store strings.

Not only do we allow accessing what Python has stored internally,
but we also support coercing this data to a `Cow<str>`.

Closes #1776.

CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 - Add `PyList::get_item_unchecked()` and `PyTuple::get_item_unchecked()` to get items without bounds checks. [#1733](https://github.com/PyO3/pyo3/pull/1733)
 - Add `PyAny::py()` as a convenience for `PyNativeType::py()`. [#1751](https://github.com/PyO3/pyo3/pull/1751)
+- Implement `PyString.data()` to access the raw bytes storing a Python string. [#1794](https://github.com/PyO3/pyo3/issues/1794)
 
 ### Changed
 

src/types/mod.rs

@@ -23,6 +23,8 @@ pub use self::num::PyLong as PyInt;
 pub use self::sequence::PySequence;
 pub use self::set::{PyFrozenSet, PySet};
 pub use self::slice::{PySlice, PySliceIndices};
+#[cfg(not(Py_LIMITED_API))]
+pub use self::string::PyStringData;
 pub use self::string::{PyString, PyString as PyUnicode};
 pub use self::tuple::PyTuple;
 pub use self::typeobject::PyType;

src/types/string.rs

@@ -1,14 +1,133 @@
 // Copyright (c) 2017-present PyO3 Project and Contributors
 
+#[cfg(not(Py_LIMITED_API))]
+use crate::exceptions::PyUnicodeDecodeError;
 use crate::types::PyBytes;
 use crate::{
     ffi, AsPyPointer, FromPyObject, IntoPy, PyAny, PyObject, PyResult, PyTryFrom, Python,
     ToPyObject,
 };
 use std::borrow::Cow;
+#[cfg(not(Py_LIMITED_API))]
+use std::ffi::CStr;
 use std::os::raw::c_char;
 use std::str;
 
+/// Represents raw data backing a Python `str`.
+///
+/// Python internally stores strings in various representations. This enumeration
+/// represents those variations.
+#[cfg(not(Py_LIMITED_API))]
+#[cfg_attr(docsrs, doc(cfg(not(Py_LIMITED_API))))]
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum PyStringData<'a> {
+    /// UCS1 representation.
+    Ucs1(&'a [u8]),
+
+    /// UCS2 representation.
+    Ucs2(&'a [u16]),
+
+    /// UCS4 representation.
+    Ucs4(&'a [u32]),
+}
+
+#[cfg(not(Py_LIMITED_API))]
+impl<'a> PyStringData<'a> {
+    /// Obtain the raw bytes backing this instance as a [u8] slice.
+    pub fn as_bytes(&self) -> &[u8] {
+        match self {
+            Self::Ucs1(s) => s,
+            Self::Ucs2(s) => unsafe {
+                std::slice::from_raw_parts(
+                    s.as_ptr() as *const u8,
+                    s.len() * self.value_width_bytes(),
+                )
+            },
+            Self::Ucs4(s) => unsafe {
+                std::slice::from_raw_parts(
+                    s.as_ptr() as *const u8,
+                    s.len() * self.value_width_bytes(),
+                )
+            },
+        }
+    }
+
+    /// Size in bytes of each value/item in the underlying slice.
+    #[inline]
+    pub fn value_width_bytes(&self) -> usize {
+        match self {
+            Self::Ucs1(_) => 1,
+            Self::Ucs2(_) => 2,
+            Self::Ucs4(_) => 4,
+        }
+    }
+
+    /// Convert the raw data to a Rust string.
+    ///
+    /// For UCS-1 / UTF-8, returns a borrow into the original slice. For UCS-2 and UCS-4,
+    /// returns an owned string.
+    ///
+    /// Returns [PyUnicodeDecodeError] if the string data isn't valid in its purported
+    /// storage format. This should only occur for strings that were created via Python
+    /// C APIs that skip input validation (like `PyUnicode_FromKindAndData`) and should
+    /// never occur for strings that were created from Python code.
+    pub fn to_string(self, py: Python) -> PyResult<Cow<'a, str>> {
+        match self {
+            Self::Ucs1(data) => match str::from_utf8(data) {
+                Ok(s) => Ok(Cow::Borrowed(s)),
+                Err(e) => Err(crate::PyErr::from_instance(PyUnicodeDecodeError::new_utf8(
+                    py, data, e,
+                )?)),
+            },
+            Self::Ucs2(data) => match String::from_utf16(data) {
+                Ok(s) => Ok(Cow::Owned(s)),
+                Err(e) => {
+                    let mut message = e.to_string().as_bytes().to_vec();
+                    message.push(0);
+
+                    Err(crate::PyErr::from_instance(PyUnicodeDecodeError::new(
+                        py,
+                        CStr::from_bytes_with_nul(b"utf-16\0").unwrap(),
+                        self.as_bytes(),
+                        0..self.as_bytes().len(),
+                        CStr::from_bytes_with_nul(&message).unwrap(),
+                    )?))
+                }
+            },
+            Self::Ucs4(data) => match data.iter().map(|&c| std::char::from_u32(c)).collect() {
+                Some(s) => Ok(Cow::Owned(s)),
+                None => Err(crate::PyErr::from_instance(PyUnicodeDecodeError::new(
+                    py,
+                    CStr::from_bytes_with_nul(b"utf-32\0").unwrap(),
+                    self.as_bytes(),
+                    0..self.as_bytes().len(),
+                    CStr::from_bytes_with_nul(b"error converting utf-32\0").unwrap(),
+                )?)),
+            },
+        }
+    }
+
+    /// Convert the raw data to a Rust string, possibly with data loss.
+    ///
+    /// Invalid code points will be replaced with `U+FFFD REPLACEMENT CHARACTER`.
+    ///
+    /// Returns a borrow into original data, when possible, or owned data otherwise.
+    ///
+    /// The return value of this function should only disagree with [Self::to_string]
+    /// when that method would error.
+    pub fn to_string_lossy(self) -> Cow<'a, str> {
+        match self {
+            Self::Ucs1(data) => String::from_utf8_lossy(data),
+            Self::Ucs2(data) => Cow::Owned(String::from_utf16_lossy(data)),
+            Self::Ucs4(data) => Cow::Owned(
+                data.iter()
+                    .map(|&c| std::char::from_u32(c).unwrap_or('\u{FFFD}'))
+                    .collect(),
+            ),
+        }
+    }
+}
+
 /// Represents a Python `string` (a Unicode string object).
 ///
 /// This type is immutable.
@@ -89,6 +208,45 @@ impl PyString {
             }
         }
     }
+
+    /// Obtains the raw data backing the Python string.
+    ///
+    /// If the Python string object was created through legacy APIs, its internal
+    /// storage format will be canonicalized before data is returned.
+    #[cfg(not(Py_LIMITED_API))]
+    #[cfg_attr(docsrs, doc(cfg(not(Py_LIMITED_API))))]
+    pub fn data(&self) -> PyResult<PyStringData<'_>> {
+        let ptr = self.as_ptr();
+
+        if cfg!(not(Py_3_12)) {
+            #[allow(deprecated)]
+            let ready = unsafe { ffi::PyUnicode_READY(ptr) };
+            if ready != 0 {
+                // Exception was created on failure.
+                return Err(crate::PyErr::api_call_failed(self.py()));
+            }
+        }
+
+        // The string should be in its canonical form after calling `PyUnicode_READY()`.
+        // And non-canonical form not possible after Python 3.12. So it should be safe
+        // to call these APIs.
+        let length = unsafe { ffi::PyUnicode_GET_LENGTH(ptr) } as usize;
+        let raw_data = unsafe { ffi::PyUnicode_DATA(ptr) };
+        let kind = unsafe { ffi::PyUnicode_KIND(ptr) };
+
+        match kind {
+            ffi::PyUnicode_1BYTE_KIND => Ok(PyStringData::Ucs1(unsafe {
+                std::slice::from_raw_parts(raw_data as *const u8, length)
+            })),
+            ffi::PyUnicode_2BYTE_KIND => Ok(PyStringData::Ucs2(unsafe {
+                std::slice::from_raw_parts(raw_data as *const u16, length)
+            })),
+            ffi::PyUnicode_4BYTE_KIND => Ok(PyStringData::Ucs4(unsafe {
+                std::slice::from_raw_parts(raw_data as *const u32, length)
+            })),
+            _ => unreachable!(),
+        }
+    }
 }
 
 /// Converts a Rust `str` to a Python object.
@@ -193,8 +351,16 @@ impl FromPyObject<'_> for char {
 #[cfg(test)]
 mod tests {
     use super::PyString;
+    #[cfg(not(Py_LIMITED_API))]
+    use super::PyStringData;
+    #[cfg(not(Py_LIMITED_API))]
+    use crate::exceptions::PyUnicodeDecodeError;
+    #[cfg(not(Py_LIMITED_API))]
+    use crate::type_object::PyTypeObject;
     use crate::Python;
     use crate::{FromPyObject, PyObject, PyTryFrom, ToPyObject};
+    #[cfg(not(Py_LIMITED_API))]
+    use std::borrow::Cow;
 
     #[test]
     fn test_non_bmp() {
@@ -297,4 +463,124 @@ mod tests {
             assert_eq!(format!("{}", s), "Hello\n");
         })
     }
+
+    #[test]
+    #[cfg(not(Py_LIMITED_API))]
+    fn test_string_data_ucs1() {
+        Python::with_gil(|py| {
+            let s = PyString::new(py, "hello, world");
+            let data = s.data().unwrap();
+
+            assert_eq!(data, PyStringData::Ucs1(b"hello, world"));
+            assert_eq!(data.to_string(py).unwrap(), Cow::Borrowed("hello, world"));
+            assert_eq!(data.to_string_lossy(), Cow::Borrowed("hello, world"));
+        })
+    }
+
+    #[test]
+    #[cfg(not(Py_LIMITED_API))]
+    fn test_string_data_ucs1_invalid() {
+        Python::with_gil(|py| {
+            // 0xfe is not allowed in UTF-8.
+            let buffer = b"f\xfe\0";
+            let ptr = unsafe {
+                crate::ffi::PyUnicode_FromKindAndData(
+                    crate::ffi::PyUnicode_1BYTE_KIND as _,
+                    buffer.as_ptr() as *const _,
+                    2,
+                )
+            };
+            assert!(!ptr.is_null());
+            let s: &PyString = unsafe { py.from_owned_ptr(ptr) };
+            let data = s.data().unwrap();
+            assert_eq!(data, PyStringData::Ucs1(b"f\xfe"));
+            let err = data.to_string(py).unwrap_err();
+            assert_eq!(err.ptype(py), PyUnicodeDecodeError::type_object(py));
+            assert!(err
+                .to_string()
+                .contains("'utf-8' codec can't decode byte 0xfe in position 1"));
+            assert_eq!(data.to_string_lossy(), Cow::Borrowed("f<EFBFBD>"));
+        });
+    }
+
+    #[test]
+    #[cfg(not(Py_LIMITED_API))]
+    fn test_string_data_ucs2() {
+        Python::with_gil(|py| {
+            let s = py.eval("'foo\\ud800'", None, None).unwrap();
+            let py_string = s.cast_as::<PyString>().unwrap();
+            let data = py_string.data().unwrap();
+
+            assert_eq!(data, PyStringData::Ucs2(&[102, 111, 111, 0xd800]));
+            assert_eq!(
+                data.to_string_lossy(),
+                Cow::Owned::<str>("foo<EFBFBD>".to_string())
+            );
+        })
+    }
+
+    #[test]
+    #[cfg(not(Py_LIMITED_API))]
+    fn test_string_data_ucs2_invalid() {
+        Python::with_gil(|py| {
+            // U+FF22 (valid) & U+d800 (never valid)
+            let buffer = b"\x22\xff\x00\xd8\x00\x00";
+            let ptr = unsafe {
+                crate::ffi::PyUnicode_FromKindAndData(
+                    crate::ffi::PyUnicode_2BYTE_KIND as _,
+                    buffer.as_ptr() as *const _,
+                    2,
+                )
+            };
+            assert!(!ptr.is_null());
+            let s: &PyString = unsafe { py.from_owned_ptr(ptr) };
+            let data = s.data().unwrap();
+            assert_eq!(data, PyStringData::Ucs2(&[0xff22, 0xd800]));
+            let err = data.to_string(py).unwrap_err();
+            assert_eq!(err.ptype(py), PyUnicodeDecodeError::type_object(py));
+            assert!(err
+                .to_string()
+                .contains("'utf-16' codec can't decode bytes in position 0-3"));
+            assert_eq!(data.to_string_lossy(), Cow::Owned::<str>("Ｂ<EFBFBD>".into()));
+        });
+    }
+
+    #[test]
+    #[cfg(not(Py_LIMITED_API))]
+    fn test_string_data_ucs4() {
+        Python::with_gil(|py| {
+            let s = "哈哈🐈";
+            let py_string = PyString::new(py, s);
+            let data = py_string.data().unwrap();
+
+            assert_eq!(data, PyStringData::Ucs4(&[21704, 21704, 128008]));
+            assert_eq!(data.to_string_lossy(), Cow::Owned::<str>(s.to_string()));
+        })
+    }
+
+    #[test]
+    #[cfg(not(Py_LIMITED_API))]
+    fn test_string_data_ucs4_invalid() {
+        Python::with_gil(|py| {
+            // U+20000 (valid) & U+d800 (never valid)
+            let buffer = b"\x00\x00\x02\x00\x00\xd8\x00\x00\x00\x00\x00\x00";
+            let ptr = unsafe {
+                crate::ffi::PyUnicode_FromKindAndData(
+                    crate::ffi::PyUnicode_4BYTE_KIND as _,
+                    buffer.as_ptr() as *const _,
+                    2,
+                )
+            };
+            assert!(!ptr.is_null());
+            let s: &PyString = unsafe { py.from_owned_ptr(ptr) };
+            let data = s.data().unwrap();
+            assert_eq!(data, PyStringData::Ucs4(&[0x20000, 0xd800]));
+            let err = data.to_string(py).unwrap_err();
+            assert_eq!(err.ptype(py), PyUnicodeDecodeError::type_object(py));
+            assert!(err
+                .to_string()
+                .contains("'utf-32' codec can't decode bytes in position 0-7"));
+            assert_eq!(data.to_string_lossy(), Cow::Owned::<str>("𠀀<EFBFBD>".into()));
+        });
+    }
 }