From 466ffeaf9f4e66542e742c82e1c916afc62cc8c4 Mon Sep 17 00:00:00 2001
From: Georg Brandl <georg@python.org>
Date: Tue, 1 Sep 2020 08:33:48 +0200
Subject: [PATCH] Avoid using CString where unnecessary

Use APIs that take a Python string instead of those taking a C string,
which are a convenience in C but require an unncessary allocation
for us.
---
 src/ffi/import.rs   |  1 +
 src/type_object.rs  | 15 +++++++--------
 src/types/mod.rs    |  1 +
 src/types/module.rs |  6 ++++--
 src/types/string.rs | 18 ++++++++++++++++++
 5 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/src/ffi/import.rs b/src/ffi/import.rs
index fdf431f1..037c54a4 100644
--- a/src/ffi/import.rs
+++ b/src/ffi/import.rs
@@ -63,6 +63,7 @@ pub unsafe fn PyImport_ImportModuleEx(
 
 extern "C" {
     pub fn PyImport_GetImporter(path: *mut PyObject) -> *mut PyObject;
+    #[cfg_attr(PyPy, link_name = "PyPyImport_Import")]
     pub fn PyImport_Import(name: *mut PyObject) -> *mut PyObject;
     #[cfg_attr(PyPy, link_name = "PyPyImport_ReloadModule")]
     pub fn PyImport_ReloadModule(m: *mut PyObject) -> *mut PyObject;
diff --git a/src/type_object.rs b/src/type_object.rs
index 546f8782..456835d8 100644
--- a/src/type_object.rs
+++ b/src/type_object.rs
@@ -229,17 +229,16 @@ fn initialize_tp_dict(
     tp_dict: *mut ffi::PyObject,
     items: Vec<(&'static str, PyObject)>,
 ) -> PyResult<()> {
-    use std::ffi::CString;
-
     // We hold the GIL: the dictionary update can be considered atomic from
     // the POV of other threads.
     for (key, val) in items {
-        let ret = unsafe {
-            ffi::PyDict_SetItemString(tp_dict, CString::new(key)?.as_ptr(), val.into_ptr())
-        };
-        if ret < 0 {
-            return Err(PyErr::fetch(py));
-        }
+        crate::types::with_tmp_string(py, key, |key| {
+            let ret = unsafe { ffi::PyDict_SetItem(tp_dict, key, val.into_ptr()) };
+            if ret < 0 {
+                return Err(PyErr::fetch(py));
+            }
+            Ok(())
+        })?;
     }
     Ok(())
 }
diff --git a/src/types/mod.rs b/src/types/mod.rs
index 32bbfe8e..abdcdd00 100644
--- a/src/types/mod.rs
+++ b/src/types/mod.rs
@@ -21,6 +21,7 @@ pub use self::num::PyLong as PyInt;
 pub use self::sequence::PySequence;
 pub use self::set::{PyFrozenSet, PySet};
 pub use self::slice::{PySlice, PySliceIndices};
+pub(crate) use self::string::with_tmp_string;
 pub use self::string::{PyString, PyString as PyUnicode};
 pub use self::tuple::PyTuple;
 pub use self::typeobject::PyType;
diff --git a/src/types/module.rs b/src/types/module.rs
index 0e480010..b77a6fa4 100644
--- a/src/types/module.rs
+++ b/src/types/module.rs
@@ -25,14 +25,16 @@ pyobject_native_var_type!(PyModule, ffi::PyModule_Type, ffi::PyModule_Check);
 impl PyModule {
     /// Creates a new module object with the `__name__` attribute set to name.
     pub fn new<'p>(py: Python<'p>, name: &str) -> PyResult<&'p PyModule> {
+        // Could use PyModule_NewObject, but it doesn't exist on PyPy.
         let name = CString::new(name)?;
         unsafe { py.from_owned_ptr_or_err(ffi::PyModule_New(name.as_ptr())) }
     }
 
     /// Imports the Python module with the specified name.
     pub fn import<'p>(py: Python<'p>, name: &str) -> PyResult<&'p PyModule> {
-        let name = CString::new(name)?;
-        unsafe { py.from_owned_ptr_or_err(ffi::PyImport_ImportModule(name.as_ptr())) }
+        crate::types::with_tmp_string(py, name, |name| unsafe {
+            py.from_owned_ptr_or_err(ffi::PyImport_Import(name))
+        })
     }
 
     /// Loads the Python code specified into a new module.
diff --git a/src/types/string.rs b/src/types/string.rs
index 722055cb..8e313edc 100644
--- a/src/types/string.rs
+++ b/src/types/string.rs
@@ -78,6 +78,24 @@ impl PyString {
     }
 }
 
+/// Convenience for calling Python APIs with a temporary string
+/// object created from a given Rust string.
+pub(crate) fn with_tmp_string<F, R>(py: Python, s: &str, f: F) -> PyResult<R>
+where
+    F: FnOnce(*mut ffi::PyObject) -> PyResult<R>,
+{
+    unsafe {
+        let s_obj =
+            ffi::PyUnicode_FromStringAndSize(s.as_ptr() as *const _, s.len() as ffi::Py_ssize_t);
+        if s_obj.is_null() {
+            return Err(PyErr::fetch(py));
+        }
+        let ret = f(s_obj);
+        ffi::Py_DECREF(s_obj);
+        ret
+    }
+}
+
 /// Converts a Rust `str` to a Python object.
 /// See `PyString::new` for details on the conversion.
 impl ToPyObject for str {