From 25a047df9b9527507c0b8b723930cd594e3b59d1 Mon Sep 17 00:00:00 2001
From: Adam Reichold <adam.reichold@t-online.de>
Date: Sun, 3 Apr 2022 10:29:35 +0200
Subject: [PATCH] Make PyString::intern more convenient by accept a string
 slice and handling null-termination internally.

---
 src/types/string.rs | 33 ++++++++++++++++++---------------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/src/types/string.rs b/src/types/string.rs
index 2b4c77da5b7..661129e1030 100644
--- a/src/types/string.rs
+++ b/src/types/string.rs
@@ -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;
 
@@ -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].
@@ -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());
         });