#[pymodule] mod some_module { ... } v3

Based on PyO3#2367 and PyO3#3294 Allows to export classes, native classes, functions and submodules and provide an init function See test/test_module.rs for an example Future work: - update examples, README and guide - investigate having #[pyclass] and #[pyfunction] directly in the #[pymodule] Co-authored-by: David Hewitt <mail@davidhewitt.dev> Co-authored-by: Georg Brandl <georg@python.org>
Tpt · Feb 21, 2024 · 50f36fa · 50f36fa
1 parent 61bc02d
commit 50f36fa
Show file tree

Hide file tree

Showing 20 changed files with 417 additions and 56 deletions.
diff --git a/newsfragments/3815.added.md b/newsfragments/3815.added.md
@@ -0,0 +1 @@
+The ability to create Python modules with a Rust `mod` block.
diff --git a/pyo3-macros-backend/src/lib.rs b/pyo3-macros-backend/src/lib.rs
@@ -22,7 +22,9 @@ mod pymethod;
 mod quotes;
 
 pub use frompyobject::build_derive_from_pyobject;
-pub use module::{process_functions_in_module, pymodule_impl, PyModuleOptions};
+pub use module::{
+    process_functions_in_module, pymodule_function_impl, pymodule_module_impl, PyModuleOptions,
+};
 pub use pyclass::{build_py_class, build_py_enum, PyClassArgs};
 pub use pyfunction::{build_py_function, PyFunctionOptions};
 pub use pyimpl::{build_py_methods, PyClassMethodsType};

diff --git a/pyo3-macros-backend/src/module.rs b/pyo3-macros-backend/src/module.rs
@@ -2,6 +2,7 @@
 
 use crate::{
     attributes::{self, take_attributes, take_pyo3_options, CrateAttribute, NameAttribute},
+    get_doc,
     pyfunction::{impl_wrap_pyfunction, PyFunctionOptions},
     utils::{get_pyo3_crate, PythonDoc},
 };
@@ -56,9 +57,156 @@ impl PyModuleOptions {
     }
 }
 
+pub fn pymodule_module_impl(mut module: syn::ItemMod) -> Result<TokenStream> {
+    let syn::ItemMod {
+        attrs,
+        vis,
+        unsafety: _,
+        ident,
+        mod_token,
+        content,
+        semi: _,
+    } = &mut module;
+    let items = if let Some((_, items)) = content {
+        items
+    } else {
+        bail_spanned!(module.span() => "`#[pymodule]` can only be used on inline modules")
+    };
+    let options = PyModuleOptions::from_attrs(attrs)?;
+    let doc = get_doc(attrs, None);
+
+    let name = options.name.unwrap_or_else(|| ident.unraw());
+    let krate = get_pyo3_crate(&options.krate);
+    let pyinit_symbol = format!("PyInit_{}", name);
+
+    let mut module_items = Vec::new();
+    let mut module_items_cfg_attrs = Vec::new();
+
+    fn extract_use_items(
+        source: &syn::UseTree,
+        cfg_attrs: &[syn::Attribute],
+        target_items: &mut Vec<syn::Ident>,
+        target_cfg_attrs: &mut Vec<Vec<syn::Attribute>>,
+    ) -> Result<()> {
+        match source {
+            syn::UseTree::Name(name) => {
+                target_items.push(name.ident.clone());
+                target_cfg_attrs.push(cfg_attrs.to_vec());
+            }
+            syn::UseTree::Path(path) => {
+                extract_use_items(&path.tree, cfg_attrs, target_items, target_cfg_attrs)?
+            }
+            syn::UseTree::Group(group) => {
+                for tree in &group.items {
+                    extract_use_items(tree, cfg_attrs, target_items, target_cfg_attrs)?
+                }
+            }
+            syn::UseTree::Glob(glob) => {
+                bail_spanned!(glob.span() => "#[pymodule] cannot import glob statements")
+            }
+            syn::UseTree::Rename(rename) => {
+                target_items.push(rename.rename.clone());
+                target_cfg_attrs.push(cfg_attrs.to_vec());
+            }
+        }
+        Ok(())
+    }
+
+    let mut pymodule_init = None;
+
+    for item in &mut *items {
+        match item {
+            syn::Item::Use(item_use) => {
+                let mut is_pyo3 = false;
+                item_use.attrs.retain(|attr| {
+                    let found = attr.path().is_ident("pyo3");
+                    is_pyo3 |= found;
+                    !found
+                });
+                if is_pyo3 {
+                    let cfg_attrs = item_use
+                        .attrs
+                        .iter()
+                        .filter(|attr| attr.path().is_ident("cfg"))
+                        .cloned()
+                        .collect::<Vec<_>>();
+                    extract_use_items(
+                        &item_use.tree,
+                        &cfg_attrs,
+                        &mut module_items,
+                        &mut module_items_cfg_attrs,
+                    )?;
+                }
+            }
+            syn::Item::Fn(item_fn) => {
+                let mut is_module_init = false;
+                item_fn.attrs.retain(|attr| {
+                    let found = attr.path().is_ident("pymodule_init");
+                    is_module_init |= found;
+                    !found
+                });
+                if is_module_init {
+                    ensure_spanned!(pymodule_init.is_none(), item_fn.span() => "only one pymodule_init may be specified");
+                    let ident = &item_fn.sig.ident;
+                    pymodule_init = Some(quote! { #ident(module)?; });
+                }
+            }
+            item => {
+                bail_spanned!(item.span() => "only 'use' statements and and pymodule_init functions are allowed in #[pymodule]")
+            }
+        }
+    }
+
+    Ok(quote! {
+        #vis #mod_token #ident {
+            #(#items)*
+
+            pub const __PYO3_NAME: &'static str = concat!(stringify!(#name), "\0");
+
+            pub(crate) struct MakeDef;
+            impl MakeDef {
+                const fn make_def() -> #krate::impl_::pymodule::ModuleDef {
+                    use #krate::impl_::pymodule as impl_;
+
+                    const INITIALIZER: impl_::ModuleInitializer = impl_::ModuleInitializer(__pyo3_pymodule);
+                    unsafe {
+                        impl_::ModuleDef::new(__PYO3_NAME, #doc, INITIALIZER)
+                    }
+                }
+            }
+
+            pub static DEF: #krate::impl_::pymodule::ModuleDef = unsafe {
+                use #krate::impl_::pymodule as impl_;
+                impl_::ModuleDef::new(concat!(stringify!(#name), "\0"), #doc, impl_::ModuleInitializer(__pyo3_pymodule))
+            };
+
+            pub fn add_to_module(module: &#krate::types::PyModule) -> #krate::PyResult<()> {
+                module.add_submodule(DEF.make_module(module.py())?.into_ref(module.py()))
+            }
+
+            pub fn __pyo3_pymodule(_py: #krate::Python, module: &#krate::types::PyModule) -> #krate::PyResult<()> {
+                use #krate::impl_::pymodule::PyAddToModule;
+                #(
+                    #(#module_items_cfg_attrs)*
+                    #module_items::add_to_module(module)?;
+                )*
+                #pymodule_init
+                Ok(())
+            }
+
+            /// This autogenerated function is called by the python interpreter when importing
+            /// the module.
+            #[export_name = #pyinit_symbol]
+            pub unsafe extern "C" fn __pyo3_init() -> *mut #krate::ffi::PyObject {
+                #krate::impl_::trampoline::module_init(|py| DEF.make_module(py))
+            }
+        }
+    })
+}
+
 /// Generates the function that is called by the python interpreter to initialize the native
 /// module
-pub fn pymodule_impl(
+pub fn pymodule_function_impl(
     fnname: &Ident,
     options: PyModuleOptions,
     doc: PythonDoc,
@@ -75,14 +223,18 @@ pub fn pymodule_impl(
         #visibility mod #fnname {
             pub(crate) struct MakeDef;
             pub static DEF: #krate::impl_::pymodule::ModuleDef = MakeDef::make_def();
-            pub const NAME: &'static str = concat!(stringify!(#name), "\0");
+            pub const __PYO3_NAME: &'static str = concat!(stringify!(#name), "\0");
 
             /// 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 {
+            pub unsafe extern "C" fn __pyo3_init() -> *mut #krate::ffi::PyObject {
                 #krate::impl_::trampoline::module_init(|py| DEF.make_module(py))
             }
+
+            pub fn add_to_module(module: &#krate::types::PyModule) -> #krate::PyResult<()> {
+                module.add_submodule(DEF.make_module(module.py())?.into_ref(module.py()))
+            }
         }
 
         // Generate the definition inside an anonymous function in the same scope as the original function -
@@ -95,7 +247,7 @@ pub fn pymodule_impl(
                 const fn make_def() -> impl_::ModuleDef {
                     const INITIALIZER: impl_::ModuleInitializer = impl_::ModuleInitializer(#fnname);
                     unsafe {
-                        impl_::ModuleDef::new(#fnname::NAME, #doc, INITIALIZER)
+                        impl_::ModuleDef::new(#fnname::__PYO3_NAME, #doc, INITIALIZER)
                     }
                 }
             }

diff --git a/pyo3-macros-backend/src/pyfunction.rs b/pyo3-macros-backend/src/pyfunction.rs
@@ -269,6 +269,10 @@ pub fn impl_wrap_pyfunction(
         #vis mod #name {
             pub(crate) struct MakeDef;
             pub const DEF: #krate::impl_::pyfunction::PyMethodDef = MakeDef::DEF;
+
+            pub fn add_to_module(module: &#krate::types::PyModule) -> #krate::PyResult<()> {
+                module.add_function(#krate::impl_::pyfunction::_wrap_pyfunction(&DEF, module)?)
+            }
         }
 
         // Generate the definition inside an anonymous function in the same scope as the original function -

diff --git a/pyo3-macros/src/lib.rs b/pyo3-macros/src/lib.rs
@@ -6,8 +6,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, PyClassArgs, PyClassMethodsType,
-    PyFunctionOptions, PyModuleOptions,
+    get_doc, process_functions_in_module, pymodule_function_impl, pymodule_module_impl,
+    PyClassArgs, PyClassMethodsType, PyFunctionOptions, PyModuleOptions,
 };
 use quote::quote;
 use syn::{parse::Nothing, parse_macro_input};
@@ -37,6 +37,12 @@ use syn::{parse::Nothing, parse_macro_input};
 pub fn pymodule(args: TokenStream, input: TokenStream) -> TokenStream {
     parse_macro_input!(args as Nothing);
 
+    if let Ok(module) = syn::parse(input.clone()) {
+        return pymodule_module_impl(module)
+            .unwrap_or_compile_error()
+            .into();
+    }
+
     let mut ast = parse_macro_input!(input as syn::ItemFn);
     let options = match PyModuleOptions::from_attrs(&mut ast.attrs) {
         Ok(options) => options,
@@ -49,7 +55,7 @@ pub fn pymodule(args: TokenStream, input: TokenStream) -> TokenStream {
 
     let doc = get_doc(&ast.attrs, None);
 
-    let expanded = pymodule_impl(&ast.sig.ident, options, doc, &ast.vis);
+    let expanded = pymodule_function_impl(&ast.sig.ident, options, doc, &ast.vis);
 
     quote!(
         #ast

diff --git a/pytests/src/lib.rs b/pytests/src/lib.rs
@@ -1,7 +1,4 @@
 use pyo3::prelude::*;
-use pyo3::types::PyDict;
-use pyo3::wrap_pymodule;
-
 pub mod awaitable;
 pub mod buf_and_str;
 pub mod comparisons;
@@ -18,43 +15,60 @@ pub mod sequence;
 pub mod subclassing;
 
 #[pymodule]
-fn pyo3_pytests(py: Python<'_>, m: &PyModule) -> PyResult<()> {
-    m.add_wrapped(wrap_pymodule!(awaitable::awaitable))?;
+mod pyo3_pytests {
+    use super::*;
+    #[pyo3]
+    use awaitable::awaitable;
+    #[pyo3]
     #[cfg(not(Py_LIMITED_API))]
-    m.add_wrapped(wrap_pymodule!(buf_and_str::buf_and_str))?;
-    m.add_wrapped(wrap_pymodule!(comparisons::comparisons))?;
+    use buf_and_str::buf_and_str;
+    #[pyo3]
+    use comparisons::comparisons;
     #[cfg(not(Py_LIMITED_API))]
-    m.add_wrapped(wrap_pymodule!(datetime::datetime))?;
-    m.add_wrapped(wrap_pymodule!(dict_iter::dict_iter))?;
-    m.add_wrapped(wrap_pymodule!(enums::enums))?;
-    m.add_wrapped(wrap_pymodule!(misc::misc))?;
-    m.add_wrapped(wrap_pymodule!(objstore::objstore))?;
-    m.add_wrapped(wrap_pymodule!(othermod::othermod))?;
-    m.add_wrapped(wrap_pymodule!(path::path))?;
-    m.add_wrapped(wrap_pymodule!(pyclasses::pyclasses))?;
-    m.add_wrapped(wrap_pymodule!(pyfunctions::pyfunctions))?;
-    m.add_wrapped(wrap_pymodule!(sequence::sequence))?;
-    m.add_wrapped(wrap_pymodule!(subclassing::subclassing))?;
-
-    // Inserting to sys.modules allows importing submodules nicely from Python
-    // e.g. import pyo3_pytests.buf_and_str as bas
-
-    let sys = PyModule::import(py, "sys")?;
-    let sys_modules: &PyDict = sys.getattr("modules")?.downcast()?;
-    sys_modules.set_item("pyo3_pytests.awaitable", m.getattr("awaitable")?)?;
-    sys_modules.set_item("pyo3_pytests.buf_and_str", m.getattr("buf_and_str")?)?;
-    sys_modules.set_item("pyo3_pytests.comparisons", m.getattr("comparisons")?)?;
-    sys_modules.set_item("pyo3_pytests.datetime", m.getattr("datetime")?)?;
-    sys_modules.set_item("pyo3_pytests.dict_iter", m.getattr("dict_iter")?)?;
-    sys_modules.set_item("pyo3_pytests.enums", m.getattr("enums")?)?;
-    sys_modules.set_item("pyo3_pytests.misc", m.getattr("misc")?)?;
-    sys_modules.set_item("pyo3_pytests.objstore", m.getattr("objstore")?)?;
-    sys_modules.set_item("pyo3_pytests.othermod", m.getattr("othermod")?)?;
-    sys_modules.set_item("pyo3_pytests.path", m.getattr("path")?)?;
-    sys_modules.set_item("pyo3_pytests.pyclasses", m.getattr("pyclasses")?)?;
-    sys_modules.set_item("pyo3_pytests.pyfunctions", m.getattr("pyfunctions")?)?;
-    sys_modules.set_item("pyo3_pytests.sequence", m.getattr("sequence")?)?;
-    sys_modules.set_item("pyo3_pytests.subclassing", m.getattr("subclassing")?)?;
+    #[pyo3]
+    use datetime::datetime;
+    #[pyo3]
+    use dict_iter::dict_iter;
+    #[pyo3]
+    use enums::enums;
+    #[pyo3]
+    use misc::misc;
+    #[pyo3]
+    use objstore::objstore;
+    #[pyo3]
+    use othermod::othermod;
+    #[pyo3]
+    use path::path;
+    #[pyo3]
+    use pyclasses::pyclasses;
+    #[pyo3]
+    use pyfunctions::pyfunctions;
+    use pyo3::types::PyDict;
+    #[pyo3]
+    use sequence::sequence;
+    #[pyo3]
+    use subclassing::subclassing;
 
-    Ok(())
+    #[pymodule_init]
+    fn init(m: &PyModule) -> PyResult<()> {
+        // Inserting to sys.modules allows importing submodules nicely from Python
+        // e.g. import pyo3_pytests.buf_and_str as bas
+        let sys = PyModule::import(m.py(), "sys")?;
+        let sys_modules: &PyDict = sys.getattr("modules")?.downcast()?;
+        sys_modules.set_item("pyo3_pytests.awaitable", m.getattr("awaitable")?)?;
+        sys_modules.set_item("pyo3_pytests.buf_and_str", m.getattr("buf_and_str")?)?;
+        sys_modules.set_item("pyo3_pytests.comparisons", m.getattr("comparisons")?)?;
+        sys_modules.set_item("pyo3_pytests.datetime", m.getattr("datetime")?)?;
+        sys_modules.set_item("pyo3_pytests.dict_iter", m.getattr("dict_iter")?)?;
+        sys_modules.set_item("pyo3_pytests.enums", m.getattr("enums")?)?;
+        sys_modules.set_item("pyo3_pytests.misc", m.getattr("misc")?)?;
+        sys_modules.set_item("pyo3_pytests.objstore", m.getattr("objstore")?)?;
+        sys_modules.set_item("pyo3_pytests.othermod", m.getattr("othermod")?)?;
+        sys_modules.set_item("pyo3_pytests.path", m.getattr("path")?)?;
+        sys_modules.set_item("pyo3_pytests.pyclasses", m.getattr("pyclasses")?)?;
+        sys_modules.set_item("pyo3_pytests.pyfunctions", m.getattr("pyfunctions")?)?;
+        sys_modules.set_item("pyo3_pytests.sequence", m.getattr("sequence")?)?;
+        sys_modules.set_item("pyo3_pytests.subclassing", m.getattr("subclassing")?)?;
+        Ok(())
+    }
 }
diff --git a/src/impl_/pymodule.rs b/src/impl_/pymodule.rs
@@ -7,7 +7,7 @@ use portable_atomic::{AtomicI64, Ordering};
 
 #[cfg(not(PyPy))]
 use crate::exceptions::PyImportError;
-use crate::{ffi, sync::GILOnceCell, types::PyModule, Py, PyResult, Python};
+use crate::{ffi, sync::GILOnceCell, types::PyModule, Py, PyResult, PyTypeInfo, Python};
 
 /// `Sync` wrapper of `ffi::PyModuleDef`.
 pub struct ModuleDef {
@@ -133,6 +133,19 @@ impl ModuleDef {
     }
 }
 
+/// Trait to add an element (class, function...) to a module.
+///
+/// Currently only implemented for classes.
+pub trait PyAddToModule {
+    fn add_to_module(module: &PyModule) -> PyResult<()>;
+}
+
+impl<T: PyTypeInfo> PyAddToModule for T {
+    fn add_to_module(module: &PyModule) -> PyResult<()> {
+        module.add(Self::NAME, Self::type_object(module.py()))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::atomic::{AtomicBool, Ordering};

diff --git a/src/macros.rs b/src/macros.rs
@@ -169,8 +169,8 @@ macro_rules! append_to_inittab {
                 );
             }
             $crate::ffi::PyImport_AppendInittab(
-                $module::NAME.as_ptr() as *const ::std::os::raw::c_char,
-                ::std::option::Option::Some($module::init),
+                $module::__PYO3_NAME.as_ptr() as *const ::std::os::raw::c_char,
+                ::std::option::Option::Some($module::__pyo3_init),
             );
         }
     };
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		The ability to create Python modules with a Rust `mod` block.