wrap_x: change macros back to macro_rules!

PyO3 · May 10, 2022 · cca0e0b · cca0e0b
1 parent 76c09ac
commit cca0e0b
Show file tree

Hide file tree

Showing 12 changed files with 102 additions and 127 deletions.
diff --git a/pyo3-macros-backend/src/lib.rs b/pyo3-macros-backend/src/lib.rs
@@ -26,7 +26,6 @@ mod pyimpl;
 mod pymethod;
 #[cfg(feature = "pyproto")]
 mod pyproto;
-mod wrap;
 
 pub use frompyobject::build_derive_from_pyobject;
 pub use module::{process_functions_in_module, pymodule_impl, PyModuleOptions};
@@ -36,4 +35,3 @@ pub use pyimpl::{build_py_methods, PyClassMethodsType};
 #[cfg(feature = "pyproto")]
 pub use pyproto::build_py_proto;
 pub use utils::get_doc;
-pub use wrap::{wrap_pyfunction_impl, wrap_pymodule_impl, WrapPyFunctionArgs};
diff --git a/pyo3-macros-backend/src/module.rs b/pyo3-macros-backend/src/module.rs
@@ -7,9 +7,8 @@ use crate::{
     },
     pyfunction::{impl_wrap_pyfunction, PyFunctionOptions},
     utils::{get_pyo3_crate, PythonDoc},
-    wrap::module_def_ident,
 };
-use proc_macro2::{Span, TokenStream};
+use proc_macro2::TokenStream;
 use quote::quote;
 use syn::{
     ext::IdentExt,
@@ -70,29 +69,49 @@ pub fn pymodule_impl(
 ) -> TokenStream {
     let name = options.name.unwrap_or_else(|| fnname.unraw());
     let krate = get_pyo3_crate(&options.krate);
-    let cb_name = Ident::new(&format!("PyInit_{}", name), Span::call_site());
-
-    let module_def_name = module_def_ident(fnname);
+    let pyinit_symbol = format!("PyInit_{}", name);
 
     quote! {
-        #[no_mangle]
-        #[allow(non_snake_case)]
-        /// This autogenerated function is called by the python interpreter when importing
-        /// the module.
-        pub unsafe extern "C" fn #cb_name() -> *mut #krate::ffi::PyObject {
-            unsafe { #module_def_name.module_init() }
+        // Create a module with the same name as the `#[pymodule]` - this way `use <the module>`
+        // will actually bring both the module and the function into scope.
+        #[doc(hidden)]
+        #visibility mod #fnname {
+            pub(crate) struct MakeDef;
+            pub static DEF: #krate::impl_::pymodule::ModuleDef = MakeDef::make_def();
+
+            /// This autogenerated function is called by the python interpreter when importing
+            /// the module.
+            #[export_name = #pyinit_symbol]
+            pub unsafe extern "C" fn init() -> *mut #krate::ffi::PyObject {
+                DEF.module_init()
+            }
         }
 
-        #[doc(hidden)]
-        #visibility static #module_def_name: #krate::impl_::pymodule::ModuleDef = unsafe {
-            #krate::impl_::pymodule::ModuleDef::new(concat!(stringify!(#name), "\0"), #doc, #krate::impl_::pymodule::ModuleInitializer(#fnname))
+        // Generate the definition inside an anonymous function in the same scope as the original function -
+        // this avoids complications around the fact that the generated module has a different scope
+        // (and `super` doesn't always refer to the outer scope, e.g. if the `#[pymodule] is
+        // inside a function body)
+        const _: () = {
+            use #krate::impl_::pymodule as impl_;
+            impl #fnname::MakeDef {
+                const fn make_def() -> impl_::ModuleDef {
+                    const INITIALIZER: impl_::ModuleInitializer = impl_::ModuleInitializer(#fnname);
+                    unsafe {
+                        impl_::ModuleDef::new(concat!(stringify!(#name), "\0"), #doc, INITIALIZER)
+                    }
+                }
+            }
         };
     }
 }
 
 /// Finds and takes care of the #[pyfn(...)] in `#[pymodule]`
-pub fn process_functions_in_module(func: &mut syn::ItemFn) -> syn::Result<()> {
+pub fn process_functions_in_module(
+    options: &PyModuleOptions,
+    func: &mut syn::ItemFn,
+) -> syn::Result<()> {
     let mut stmts: Vec<syn::Stmt> = Vec::new();
+    let krate = get_pyo3_crate(&options.krate);
 
     for mut stmt in func.block.stmts.drain(..) {
         if let syn::Stmt::Item(syn::Item::Fn(func)) = &mut stmt {
@@ -102,7 +121,7 @@ pub fn process_functions_in_module(func: &mut syn::ItemFn) -> syn::Result<()> {
                 let name = &func.sig.ident;
                 let statements: Vec<syn::Stmt> = syn::parse_quote! {
                     #wrapped_function
-                    #module_name.add_function(#name::wrap(#name::DEF, #module_name)?)?;
+                    #module_name.add_function(#krate::impl_::pyfunction::wrap_pyfunction(&#name::DEF, #module_name)?)?;
                 };
                 stmts.extend(statements);
             }

diff --git a/pyo3-macros-backend/src/pyfunction.rs b/pyo3-macros-backend/src/pyfunction.rs
@@ -446,12 +446,8 @@ pub fn impl_wrap_pyfunction(
         // will actually bring both the module and the function into scope.
         #[doc(hidden)]
         #vis mod #name {
-            use #krate as _pyo3;
-            pub(crate) struct PyO3Def;
-
-            // Exported for `wrap_pyfunction!`
-            pub use _pyo3::impl_::pyfunction::wrap_pyfunction as wrap;
-            pub const DEF: _pyo3::PyMethodDef = <PyO3Def as _pyo3::impl_::pyfunction::PyFunctionDef>::DEF;
+            pub(crate) struct MakeDef;
+            pub const DEF: #krate::impl_::pyfunction::PyMethodDef = MakeDef::DEF;
         }
 
         // Generate the definition inside an anonymous function in the same scope as the original function -
@@ -460,8 +456,8 @@ pub fn impl_wrap_pyfunction(
         // inside a function body)
         const _: () = {
             use #krate as _pyo3;
-            impl _pyo3::impl_::pyfunction::PyFunctionDef for #name::PyO3Def {
-                const DEF: _pyo3::PyMethodDef = #methoddef;
+            impl #name::MakeDef {
+                const DEF: #krate::impl_::pyfunction::PyMethodDef = #methoddef;
             }
         };
     };

diff --git a/pyo3-macros-backend/src/wrap.rs b/pyo3-macros-backend/src/wrap.rs
diff --git a/pyo3-macros/src/lib.rs b/pyo3-macros/src/lib.rs
@@ -9,8 +9,8 @@ use proc_macro::TokenStream;
 use proc_macro2::TokenStream as TokenStream2;
 use pyo3_macros_backend::{
     build_derive_from_pyobject, build_py_class, build_py_enum, build_py_function, build_py_methods,
-    get_doc, process_functions_in_module, pymodule_impl, wrap_pyfunction_impl, wrap_pymodule_impl,
-    PyClassArgs, PyClassMethodsType, PyFunctionOptions, PyModuleOptions, WrapPyFunctionArgs,
+    get_doc, process_functions_in_module, pymodule_impl, PyClassArgs, PyClassMethodsType,
+    PyFunctionOptions, PyModuleOptions,
 };
 use quote::quote;
 use syn::{parse::Nothing, parse_macro_input};
@@ -41,7 +41,7 @@ pub fn pymodule(args: TokenStream, input: TokenStream) -> TokenStream {
         Err(e) => return e.into_compile_error().into(),
     };
 
-    if let Err(err) = process_functions_in_module(&mut ast) {
+    if let Err(err) = process_functions_in_module(&options, &mut ast) {
         return err.into_compile_error().into();
     }
 
@@ -194,25 +194,6 @@ pub fn derive_from_py_object(item: TokenStream) -> TokenStream {
     .into()
 }
 
-/// Wraps a Rust function annotated with [`#[pyfunction]`](macro@crate::pyfunction).
-///
-/// This can be used with `PyModule::add_function` to add free functions to a `PyModule` - see its
-/// documentation for more information.
-#[proc_macro]
-pub fn wrap_pyfunction(input: TokenStream) -> TokenStream {
-    let args = parse_macro_input!(input as WrapPyFunctionArgs);
-    wrap_pyfunction_impl(args).into()
-}
-
-/// Returns a function that takes a `Python` instance and returns a Python module.
-///
-/// Use this together with [`#[pymodule]`](macro@crate::pymodule) and `PyModule::add_wrapped`.
-#[proc_macro]
-pub fn wrap_pymodule(input: TokenStream) -> TokenStream {
-    let path = parse_macro_input!(input as syn::Path);
-    wrap_pymodule_impl(path).unwrap_or_compile_error().into()
-}
-
 fn pyclass_impl(
     attrs: TokenStream,
     mut ast: syn::ItemStruct,

diff --git a/src/impl_/pyfunction.rs b/src/impl_/pyfunction.rs
@@ -1,14 +1,10 @@
-use crate::{
-    derive_utils::PyFunctionArguments, impl_::pymethods::PyMethodDef, types::PyCFunction, PyResult,
-};
+use crate::{derive_utils::PyFunctionArguments, types::PyCFunction, PyResult};
 
-pub trait PyFunctionDef {
-    const DEF: crate::PyMethodDef;
-}
+pub use crate::impl_::pymethods::PyMethodDef;
 
 pub fn wrap_pyfunction<'a>(
-    method_def: PyMethodDef,
-    args: impl Into<PyFunctionArguments<'a>>,
+    method_def: &PyMethodDef,
+    py_or_module: impl Into<PyFunctionArguments<'a>>,
 ) -> PyResult<&'a PyCFunction> {
-    PyCFunction::internal_new(method_def, args.into())
+    PyCFunction::internal_new(method_def, py_or_module.into())
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -405,9 +405,7 @@ pub mod proc_macro {
 #[cfg(all(feature = "macros", feature = "pyproto"))]
 pub use pyo3_macros::pyproto;
 #[cfg(feature = "macros")]
-pub use pyo3_macros::{
-    pyclass, pyfunction, pymethods, pymodule, wrap_pyfunction, wrap_pymodule, FromPyObject,
-};
+pub use pyo3_macros::{pyclass, pyfunction, pymethods, pymodule, FromPyObject};
 
 #[cfg(feature = "macros")]
 #[macro_use]

diff --git a/src/macros.rs b/src/macros.rs
@@ -115,3 +115,36 @@ macro_rules! py_run_impl {
         }
     }};
 }
+
+/// Wraps a Rust function annotated with [`#[pyfunction]`](macro@crate::pyfunction).
+///
+/// This can be used with [`PyModule::add_function`](crate::types::PyModule::add_function) to add free
+/// functions to a [`PyModule`](crate::types::PyModule) - see its documentation for more information.
+#[macro_export]
+macro_rules! wrap_pyfunction {
+    ($function:path) => {
+        &|py_or_module| {
+            use $function as def;
+            $crate::impl_::pyfunction::wrap_pyfunction(&def::DEF, py_or_module)
+        }
+    };
+    ($function:path, $py_or_module:expr) => {{
+        use $function as def;
+        $crate::impl_::pyfunction::wrap_pyfunction(&def::DEF, $py_or_module)
+    }};
+}
+
+/// Returns a function that takes a [`Python`](crate::Python) instance and returns a
+/// Python module.
+///
+/// Use this together with [`#[pymodule]`](crate::pymodule) and
+/// [`PyModule::add_wrapped`](crate::types::PyModule::add_wrapped).
+#[macro_export]
+macro_rules! wrap_pymodule {
+    ($module:path) => {
+        &|py| {
+            use $module as def;
+            def::DEF.make_module(py).expect("failed to wrap pymodule")
+        }
+    };
+}
diff --git a/src/once_cell.rs b/src/once_cell.rs
@@ -140,6 +140,9 @@ impl<T> GILOnceCell<T> {
 /// }
 /// #
 /// # Python::with_gil(|py| {
+/// #     let fun_slow = wrap_pyfunction!(create_dict, py).unwrap();
+/// #     let dict = fun.call0().unwrap();
+/// #     assert!(dict.contains("foo").unwrap());
 /// #     let fun = wrap_pyfunction!(create_dict_faster, py).unwrap();
 /// #     let dict = fun.call0().unwrap();
 /// #     assert!(dict.contains("foo").unwrap());

diff --git a/src/prelude.rs b/src/prelude.rs
@@ -22,7 +22,10 @@ pub use crate::pyclass_init::PyClassInitializer;
 pub use crate::types::{PyAny, PyModule};
 
 #[cfg(feature = "macros")]
-pub use pyo3_macros::{pyclass, pyfunction, pymethods, pymodule, wrap_pyfunction, FromPyObject};
+pub use pyo3_macros::{pyclass, pyfunction, pymethods, pymodule, FromPyObject};
 
 #[cfg(all(feature = "macros", feature = "pyproto"))]
 pub use pyo3_macros::pyproto;
+
+#[cfg(feature = "macros")]
+pub use crate::wrap_pyfunction;
diff --git a/src/types/bytearray.rs b/src/types/bytearray.rs
@@ -166,10 +166,11 @@ impl PyByteArray {
     ///
     /// The following `bug` function is unsound ⚠️
     ///
-    /// ```rust
+    /// ```rust,no_run
     /// # use pyo3::prelude::*;
     /// # use pyo3::types::PyByteArray;
     ///
+    /// # #[allow(dead_code)]
     /// #[pyfunction]
     /// fn bug(py: Python<'_>, bytes: &PyByteArray) {
     ///     let slice = unsafe { bytes.as_bytes() };
@@ -186,6 +187,7 @@ impl PyByteArray {
     ///     // remaining valid. As such this is also undefined behavior.
     ///     println!("{:?}", slice[0]);
     /// }
+    /// ```
     pub unsafe fn as_bytes(&self) -> &[u8] {
         slice::from_raw_parts(self.data(), self.len())
     }

diff --git a/src/types/function.rs b/src/types/function.rs
@@ -65,7 +65,7 @@ impl PyCFunction {
         py_or_module: PyFunctionArguments<'a>,
     ) -> PyResult<&'a Self> {
         Self::internal_new(
-            PyMethodDef::cfunction_with_keywords(
+            &PyMethodDef::cfunction_with_keywords(
                 name,
                 pymethods::PyCFunctionWithKeywords(fun),
                 doc,
@@ -82,7 +82,7 @@ impl PyCFunction {
         py_or_module: PyFunctionArguments<'a>,
     ) -> PyResult<&'a Self> {
         Self::internal_new(
-            PyMethodDef::noargs(name, pymethods::PyCFunction(fun), doc),
+            &PyMethodDef::noargs(name, pymethods::PyCFunction(fun), doc),
             py_or_module,
         )
     }
@@ -125,16 +125,16 @@ impl PyCFunction {
             pymethods::PyCFunctionWithKeywords(run_closure::<F, R>),
             "",
         );
-        Self::internal_new_from_pointers(method_def, py, capsule.as_ptr(), std::ptr::null_mut())
+        Self::internal_new_from_pointers(&method_def, py, capsule.as_ptr(), std::ptr::null_mut())
     }
 
     #[doc(hidden)]
-    fn internal_new_from_pointers(
-        method_def: PyMethodDef,
-        py: Python<'_>,
+    fn internal_new_from_pointers<'py>(
+        method_def: &PyMethodDef,
+        py: Python<'py>,
         mod_ptr: *mut ffi::PyObject,
         module_name: *mut ffi::PyObject,
-    ) -> PyResult<&Self> {
+    ) -> PyResult<&'py Self> {
         let def = method_def
             .as_method_def()
             .map_err(|err| PyValueError::new_err(err.0))?;
@@ -148,10 +148,10 @@ impl PyCFunction {
     }
 
     #[doc(hidden)]
-    pub fn internal_new(
-        method_def: PyMethodDef,
-        py_or_module: PyFunctionArguments<'_>,
-    ) -> PyResult<&Self> {
+    pub fn internal_new<'py>(
+        method_def: &PyMethodDef,
+        py_or_module: PyFunctionArguments<'py>,
+    ) -> PyResult<&'py Self> {
         let (py, module) = py_or_module.into_py_and_maybe_module();
         let (mod_ptr, module_name) = if let Some(m) = module {
             let mod_ptr = m.as_ptr();