Skip to content

Commit

Permalink
Make PyString::intern more convenient by accept a string slice and ha…
Browse files Browse the repository at this point in the history
…ndling null-termination internally.
  • Loading branch information
adamreichold committed Apr 3, 2022
1 parent d539d24 commit 25a047d
Showing 1 changed file with 18 additions and 15 deletions.
33 changes: 18 additions & 15 deletions src/types/string.rs
Expand Up @@ -8,7 +8,7 @@ use crate::{
ToPyObject,
};
use std::borrow::Cow;
use std::ffi::CStr;
use std::ffi::{CStr, CString};
use std::os::raw::c_char;
use std::str;

Expand Down Expand Up @@ -144,26 +144,31 @@ impl PyString {
unsafe { py.from_owned_ptr(ffi::PyUnicode_FromStringAndSize(ptr, len)) }
}

/// Intern the given null-terminated string
/// Intern the given string
///
/// This will return a reference to the same Python string object if called repeatedly with the same string.
///
/// Panics if out of memory.
/// Note that a temporary allocation is required if the given string is not already null-terminated.
///
/// Panics if out of memory or null byte within string.
///
/// # Example
///
/// ```
/// use std::ffi::CStr;
/// # use pyo3::{types::PyString, Python};
///
/// # Python::with_gil(|py| {
/// let c_str = CStr::from_bytes_with_nul(b"foobar\0").unwrap();
/// let py_str = PyString::intern(py, c_str);
/// assert_eq!(c_str.to_str().unwrap(), py_str.to_str().unwrap());
/// let py_str = PyString::intern(py, "foobar\0");
/// assert_eq!(py_str.to_str().unwrap(), "foobar");
/// # });
/// ```
pub fn intern<'p>(py: Python<'p>, s: &CStr) -> &'p PyString {
unsafe { py.from_owned_ptr(ffi::PyUnicode_InternFromString(s.as_ptr())) }
pub fn intern<'p>(py: Python<'p>, s: &str) -> &'p PyString {
if let Ok(s) = CStr::from_bytes_with_nul(s.as_bytes()) {
unsafe { py.from_owned_ptr(ffi::PyUnicode_InternFromString(s.as_ptr())) }
} else {
let s = CString::new(s).unwrap();
unsafe { py.from_owned_ptr(ffi::PyUnicode_InternFromString(s.as_ptr())) }
}
}

/// Attempts to create a Python string from a Python [bytes-like object].
Expand Down Expand Up @@ -618,13 +623,11 @@ mod tests {
#[test]
fn test_intern_string() {
Python::with_gil(|py| {
let s = CStr::from_bytes_with_nul(b"foobar\0").unwrap();

let py_string1 = PyString::intern(py, s);
assert_eq!(s.to_str().unwrap(), py_string1.to_str().unwrap());
let py_string1 = PyString::intern(py, "foobar\0");
assert_eq!(py_string1.to_str().unwrap(), "foobar");

let py_string2 = PyString::intern(py, s);
assert_eq!(s.to_str().unwrap(), py_string2.to_str().unwrap());
let py_string2 = PyString::intern(py, "foobar");
assert_eq!(py_string2.to_str().unwrap(), "foobar");

assert_eq!(py_string1.as_ptr(), py_string2.as_ptr());
});
Expand Down

0 comments on commit 25a047d

Please sign in to comment.