Better code generation

PyO3 · Sep 17, 2018 · 2904291 · 2904291
1 parent 0372360
commit 2904291
Show file tree

Hide file tree

Showing 5 changed files with 96 additions and 72 deletions.
diff --git a/pyo3-derive-backend/src/module.rs b/pyo3-derive-backend/src/module.rs
@@ -20,37 +20,7 @@ pub fn py3_init(fnname: &syn::Ident, name: &syn::Ident, doc: syn::Lit) -> TokenS
         /// This autogenerated function is called by the python interpreter when importing
         /// the module.
         pub unsafe extern "C" fn #cb_name() -> *mut ::pyo3::ffi::PyObject {
-            ::pyo3::init_once();
-
-            static mut MODULE_DEF: ::pyo3::ffi::PyModuleDef = ::pyo3::ffi::PyModuleDef_INIT;
-            // We can't convert &'static str to *const c_char within a static initializer,
-            // so we'll do it here in the module initialization:
-            MODULE_DEF.m_name = concat!(stringify!(#name), "\0").as_ptr() as *const _;
-
-            ::pyo3::PyEval_InitThreads_if_with_thread();
-
-            let _module = ::pyo3::ffi::PyModule_Create(&mut MODULE_DEF);
-            if _module.is_null() {
-                return _module;
-            }
-
-            let _pool = ::pyo3::GILPool::new();
-            let _py = ::pyo3::Python::assume_gil_acquired();
-            let _module = match _py.from_owned_ptr_or_err::<::pyo3::PyModule>(_module) {
-                Ok(m) => m,
-                Err(e) => {
-                    ::pyo3::PyErr::from(e).restore(_py);
-                    return ::std::ptr::null_mut();
-                }
-            };
-            _module.add("__doc__", #doc).expect("Failed to add doc for module");
-            match #fnname(_py, _module) {
-                Ok(_) => ::pyo3::IntoPyPointer::into_ptr(_module),
-                Err(e) => {
-                    e.restore(_py);
-                    ::std::ptr::null_mut()
-                }
-            }
+            ::pyo3::make_module(concat!(stringify!(#name), "\0"), #doc, #fnname)
         }
     }
 }
@@ -62,30 +32,7 @@ pub fn py2_init(fnname: &syn::Ident, name: &syn::Ident, doc: syn::Lit) -> TokenS
         #[no_mangle]
         #[allow(non_snake_case)]
         pub unsafe extern "C" fn #cb_name() {
-            // initialize python
-            ::pyo3::init_once();
-            ::pyo3::PyEval_InitThreads_if_with_thread();
-
-            let _name = concat!(stringify!(#name), "\0").as_ptr() as *const _;
-            let _pool = ::pyo3::GILPool::new();
-            let _py = ::pyo3::Python::assume_gil_acquired();
-            let _module = ::pyo3::ffi::Py_InitModule(_name, ::std::ptr::null_mut());
-            if _module.is_null() {
-                return
-            }
-
-            let _module = match _py.from_borrowed_ptr_or_err::<::pyo3::PyModule>(_module) {
-                Ok(m) => m,
-                Err(e) => {
-                    ::pyo3::PyErr::from(e).restore(_py);
-                    return
-                }
-            };
-
-            _module.add("__doc__", #doc).expect("Failed to add doc for module");
-            if let Err(e) = #fnname(_py, _module) {
-                e.restore(_py)
-            }
+            ::pyo3::make_module(concat!(stringify!(#name), "\0"), #doc, #fnname)
         }
     }
 }

diff --git a/src/lib.rs b/src/lib.rs
@@ -149,7 +149,7 @@ pub use object::PyObject;
 pub use objectprotocol::ObjectProtocol;
 pub use objects::*;
 pub use python::{IntoPyPointer, Python, ToPyPointer};
-pub use pythonrun::{init_once, prepare_freethreaded_python, GILGuard, GILPool, PyEval_InitThreads_if_with_thread};
+pub use pythonrun::{init_once, prepare_freethreaded_python, GILGuard, GILPool};
 pub use typeob::{PyObjectAlloc, PyRawObject, PyTypeInfo};
 pub mod class;
 pub use class::*;

diff --git a/src/objects/mod.rs b/src/objects/mod.rs
@@ -10,7 +10,7 @@ pub use self::dict::PyDict;
 pub use self::floatob::PyFloat;
 pub use self::iterator::PyIterator;
 pub use self::list::PyList;
-pub use self::module::PyModule;
+pub use self::module::{make_module, PyModule};
 #[cfg(not(Py_3))]
 pub use self::num2::{PyInt, PyLong};
 #[cfg(Py_3)]

diff --git a/src/objects/module.rs b/src/objects/module.rs
@@ -2,19 +2,21 @@
 //
 // based on Daniel Grunwald's https://github.com/dgrunwald/rust-cpython
 
-use std::ffi::{CStr, CString};
-use std::os::raw::c_char;
-
 use conversion::{IntoPyTuple, ToPyObject};
 use err::{PyErr, PyResult};
 use ffi;
+use init_once;
 use instance::PyObjectWithToken;
 use object::PyObject;
 use objectprotocol::ObjectProtocol;
 use objects::{exc, PyDict, PyObjectRef, PyType};
 use python::{Python, ToPyPointer};
+use std::ffi::{CStr, CString};
+use std::os::raw::c_char;
+use std::ptr;
 use std::str;
 use typeob::{initialize_type, PyTypeInfo};
+use GILPool;
 
 /// Represents a Python `module` object.
 #[repr(transparent)]
@@ -192,3 +194,90 @@ impl PyModule {
         self.add(name.extract(self.py()).unwrap(), function)
     }
 }
+
+#[cfg(Py_3)]
+#[doc(hidden)]
+/// Builds a module (or null) from a user given initializer. Used for `#[pymodinit]`.
+pub unsafe fn make_module(
+    name: &str,
+    doc: &str,
+    initializer: impl Fn(Python, &PyModule) -> PyResult<()>,
+) -> *mut ffi::PyObject {
+    use python::IntoPyPointer;
+
+    init_once();
+
+    #[cfg(py_sys_config = "WITH_THREAD")]
+    // > Changed in version 3.7: This function is now called by Py_Initialize(), so you don’t have
+    // > to call it yourself anymore.
+    #[cfg(not(Py_3_7))]
+    ffi::PyEval_InitThreads();
+
+    static mut MODULE_DEF: ffi::PyModuleDef = ffi::PyModuleDef_INIT;
+    // We can't convert &'static str to *const c_char within a static initializer,
+    // so we'll do it here in the module initialization:
+    MODULE_DEF.m_name = name.as_ptr() as *const _;
+
+    let module = ffi::PyModule_Create(&mut MODULE_DEF);
+    if module.is_null() {
+        return module;
+    }
+
+    let _pool = GILPool::new();
+    let py = Python::assume_gil_acquired();
+    let module = match py.from_owned_ptr_or_err::<PyModule>(module) {
+        Ok(m) => m,
+        Err(e) => {
+            e.restore(py);
+            return ptr::null_mut();
+        }
+    };
+
+    module
+        .add("__doc__", doc)
+        .expect("Failed to add doc for module");
+    match initializer(py, module) {
+        Ok(_) => module.into_ptr(),
+        Err(e) => {
+            e.restore(py);
+            ptr::null_mut()
+        }
+    }
+}
+
+#[cfg(not(Py_3))]
+#[doc(hidden)]
+/// Builds a module (or null) from a user given initializer. Used for `#[pymodinit]`.
+pub unsafe fn make_module(
+    name: &str,
+    doc: &str,
+    initializer: impl Fn(Python, &PyModule) -> PyResult<()>,
+) {
+    init_once();
+
+    #[cfg(py_sys_config = "WITH_THREAD")]
+    ffi::PyEval_InitThreads();
+
+    let _name = name.as_ptr() as *const _;
+    let _pool = GILPool::new();
+    let py = Python::assume_gil_acquired();
+    let _module = ffi::Py_InitModule(_name, ptr::null_mut());
+    if _module.is_null() {
+        return;
+    }
+
+    let _module = match py.from_borrowed_ptr_or_err::<PyModule>(_module) {
+        Ok(m) => m,
+        Err(e) => {
+            e.restore(py);
+            return;
+        }
+    };
+
+    _module
+        .add("__doc__", doc)
+        .expect("Failed to add doc for module");
+    if let Err(e) = initializer(py, _module) {
+        e.restore(py)
+    }
+}
diff --git a/src/pythonrun.rs b/src/pythonrun.rs
@@ -64,18 +64,6 @@ pub fn prepare_freethreaded_python() {
     });
 }
 
-/// This function is called in the generated module init code. We need to declare it the pyo3 crate
-/// because the user's crate won't have the cfg's.
-#[doc(hidden)]
-#[allow(non_snake_case)]
-pub unsafe fn PyEval_InitThreads_if_with_thread() {
-    #[cfg(py_sys_config = "WITH_THREAD")]
-    // > Changed in version 3.7: This function is now called by Py_Initialize(), so you don’t have
-    // > to call it yourself anymore.
-    #[cfg(not(Py_3_7))]
-    ffi::PyEval_InitThreads();
-}
-
 #[doc(hidden)]
 pub fn init_once() {
     START_PYO3.call_once(|| unsafe {