#[pymodule] mod some_module { ... } v3 (#3815)

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

Based on #2367 and #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>

* tests: group exported imports

* Consolidate pymodule macro code to avoid duplicates

* Makes pymodule_init take Bound<'_, PyModule>

* Renames #[pyo3] to #[pymodule_export]

* Gates #[pymodule] mod behind the experimental-declarative-modules feature

* Properly fails on functions inside of declarative modules

---------

Co-authored-by: David Hewitt <mail@davidhewitt.dev>
Co-authored-by: Georg Brandl <georg@python.org>

Cargo.toml

@@ -69,6 +69,9 @@ default = ["macros"]
 # and IntoPy traits
 experimental-inspect = []
 
+# Enables annotating Rust inline modules with #[pymodule] to build Python modules declaratively
+experimental-declarative-modules = ["pyo3-macros/experimental-declarative-modules", "macros"]
+
 # Enables macros: #[pyclass], #[pymodule], #[pyfunction] etc.
 macros = ["pyo3-macros", "indoc", "unindent"]
 
@@ -114,6 +117,7 @@ full = [
     "chrono-tz",
     "either",
     "experimental-inspect",
+    "experimental-declarative-modules",
     "eyre",
     "hashbrown",
     "indexmap",

newsfragments/3815.added.md

@@ -0,0 +1,2 @@
+The ability to create Python modules with a Rust `mod` block
+behind the `experimental-declarative-modules` feature.

pyo3-macros-backend/src/lib.rs

@@ -22,7 +22,7 @@ mod pymethod;
 mod quotes;
 
 pub use frompyobject::build_derive_from_pyobject;
-pub use module::{process_functions_in_module, pymodule_impl, PyModuleOptions};
+pub use 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};

pyo3-macros-backend/src/module.rs

@@ -2,8 +2,9 @@
 
 use crate::{
     attributes::{self, take_attributes, take_pyo3_options, CrateAttribute, NameAttribute},
+    get_doc,
     pyfunction::{impl_wrap_pyfunction, PyFunctionOptions},
-    utils::{get_pyo3_crate, PythonDoc},
+    utils::get_pyo3_crate,
 };
 use proc_macro2::TokenStream;
 use quote::quote;
@@ -12,7 +13,7 @@ use syn::{
     parse::{Parse, ParseStream},
     spanned::Spanned,
     token::Comma,
-    Ident, Path, Result, Visibility,
+    Item, Path, Result,
 };
 
 #[derive(Default)]
@@ -56,33 +57,154 @@ 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 krate = get_pyo3_crate(&options.krate);
+    let doc = get_doc(attrs, None);
+
+    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 {
+            Item::Use(item_use) => {
+                let mut is_pyo3 = false;
+                item_use.attrs.retain(|attr| {
+                    let found = attr.path().is_ident("pymodule_export");
+                    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,
+                    )?;
+                }
+            }
+            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)?; });
+                } else {
+                    bail_spanned!(item.span() => "only 'use' statements and and pymodule_init functions are allowed in #[pymodule]")
+                }
+            }
+            item => {
+                bail_spanned!(item.span() => "only 'use' statements and and pymodule_init functions are allowed in #[pymodule]")
+            }
+        }
+    }
+
+    let initialization = module_initialization(options, ident);
+    Ok(quote!(
+        #vis mod #ident {
+            #(#items)*
+
+            #initialization
+
+            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
+                        )
+                    }
+                }
+            }
+
+            fn __pyo3_pymodule(module: &#krate::Bound<'_, #krate::types::PyModule>) -> #krate::PyResult<()> {
+                use #krate::impl_::pymodule::PyAddToModule;
+                #(
+                    #(#module_items_cfg_attrs)*
+                    #module_items::add_to_module(module)?;
+                )*
+                #pymodule_init
+                Ok(())
+            }
+        }
+    ))
+}
+
 /// Generates the function that is called by the python interpreter to initialize the native
 /// module
-pub fn pymodule_impl(
-    fnname: &Ident,
-    options: PyModuleOptions,
-    doc: PythonDoc,
-    visibility: &Visibility,
-) -> TokenStream {
-    let name = options.name.unwrap_or_else(|| fnname.unraw());
+pub fn pymodule_function_impl(mut function: syn::ItemFn) -> Result<TokenStream> {
+    let options = PyModuleOptions::from_attrs(&mut function.attrs)?;
+    process_functions_in_module(&options, &mut function)?;
     let krate = get_pyo3_crate(&options.krate);
-    let pyinit_symbol = format!("PyInit_{}", name);
+    let ident = &function.sig.ident;
+    let vis = &function.vis;
+    let doc = get_doc(&function.attrs, None);
 
-    quote! {
-        // 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();
-            pub const 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 {
-                #krate::impl_::trampoline::module_init(|py| DEF.make_module(py))
-            }
+    let initialization = module_initialization(options, ident);
+    Ok(quote! {
+        #function
+        #vis mod #ident {
+            #initialization
         }
 
         // Generate the definition inside an anonymous function in the same scope as the original function -
@@ -91,28 +213,59 @@ pub fn pymodule_impl(
         // inside a function body)
         const _: () = {
             use #krate::impl_::pymodule as impl_;
-            impl #fnname::MakeDef {
+
+            fn __pyo3_pymodule(module: &#krate::Bound<'_, #krate::types::PyModule>) -> #krate::PyResult<()> {
+                #ident(module.py(), module.as_gil_ref())
+            }
+
+            impl #ident::MakeDef {
                 const fn make_def() -> impl_::ModuleDef {
-                    const INITIALIZER: impl_::ModuleInitializer = impl_::ModuleInitializer(#fnname);
                     unsafe {
-                        impl_::ModuleDef::new(#fnname::NAME, #doc, INITIALIZER)
+                        const INITIALIZER: impl_::ModuleInitializer = impl_::ModuleInitializer(__pyo3_pymodule);
+                        impl_::ModuleDef::new(
+                            #ident::__PYO3_NAME,
+                            #doc,
+                            INITIALIZER
+                        )
                     }
                 }
             }
         };
+    })
+}
+
+fn module_initialization(options: PyModuleOptions, ident: &syn::Ident) -> TokenStream {
+    let name = options.name.unwrap_or_else(|| ident.unraw());
+    let krate = get_pyo3_crate(&options.krate);
+    let pyinit_symbol = format!("PyInit_{}", name);
+
+    quote! {
+        pub const __PYO3_NAME: &'static str = concat!(stringify!(#name), "\0");
+
+        pub(super) struct MakeDef;
+        pub static DEF: #krate::impl_::pymodule::ModuleDef = MakeDef::make_def();
+
+        pub fn add_to_module(module: &#krate::Bound<'_, #krate::types::PyModule>) -> #krate::PyResult<()> {
+            use #krate::prelude::PyModuleMethods;
+            module.add_submodule(DEF.make_module(module.py())?.bind(module.py()))
+        }
+
+        /// 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))
+        }
     }
 }
 
 /// Finds and takes care of the #[pyfn(...)] in `#[pymodule]`
-pub fn process_functions_in_module(
-    options: &PyModuleOptions,
-    func: &mut syn::ItemFn,
-) -> syn::Result<()> {
+fn process_functions_in_module(options: &PyModuleOptions, func: &mut syn::ItemFn) -> 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 {
+        if let syn::Stmt::Item(Item::Fn(func)) = &mut stmt {
             if let Some(pyfn_args) = get_pyfn_attr(&mut func.attrs)? {
                 let module_name = pyfn_args.modname;
                 let wrapped_function = impl_wrap_pyfunction(func, pyfn_args.options)?;

pyo3-macros-backend/src/pyfunction.rs

@@ -269,6 +269,12 @@ 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::Bound<'_, #krate::types::PyModule>) -> #krate::PyResult<()> {
+                use #krate::prelude::PyModuleMethods;
+                use ::std::convert::Into;
+                module.add_function(&#krate::types::PyCFunction::internal_new(&DEF, module.as_gil_ref().into())?)
+            }
         }
 
         // Generate the definition inside an anonymous function in the same scope as the original function -

pyo3-macros/Cargo.toml

@@ -15,6 +15,7 @@ proc-macro = true
 
 [features]
 multiple-pymethods = []
+experimental-declarative-modules = []
 
 [dependencies]
 proc-macro2 = { version = "1", default-features = false }

pyo3-macros/src/lib.rs

@@ -6,11 +6,11 @@ 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,
+    pymodule_function_impl, pymodule_module_impl, PyClassArgs, PyClassMethodsType,
+    PyFunctionOptions,
 };
 use quote::quote;
-use syn::{parse::Nothing, parse_macro_input};
+use syn::{parse::Nothing, parse_macro_input, Item};
 
 /// A proc macro used to implement Python modules.
 ///
@@ -36,31 +36,27 @@ use syn::{parse::Nothing, parse_macro_input};
 #[proc_macro_attribute]
 pub fn pymodule(args: TokenStream, input: TokenStream) -> TokenStream {
     parse_macro_input!(args as Nothing);
-
-    let mut ast = parse_macro_input!(input as syn::ItemFn);
-    let options = match PyModuleOptions::from_attrs(&mut ast.attrs) {
-        Ok(options) => options,
-        Err(e) => return e.into_compile_error().into(),
-    };
-
-    if let Err(err) = process_functions_in_module(&options, &mut ast) {
-        return err.into_compile_error().into();
+    match parse_macro_input!(input as Item) {
+        Item::Mod(module) => if cfg!(feature = "experimental-declarative-modules") {
+            pymodule_module_impl(module)
+        } else {
+            Err(syn::Error::new_spanned(
+                module,
+                "#[pymodule] requires the 'experimental-declarative-modules' feature to be used on Rust modules.",
+            ))
+        },
+        Item::Fn(function) => pymodule_function_impl(function),
+        unsupported => Err(syn::Error::new_spanned(
+            unsupported,
+            "#[pymodule] only supports modules and functions.",
+        )),
     }
-
-    let doc = get_doc(&ast.attrs, None);
-
-    let expanded = pymodule_impl(&ast.sig.ident, options, doc, &ast.vis);
-
-    quote!(
-        #ast
-        #expanded
-    )
+    .unwrap_or_compile_error()
     .into()
 }
 
 #[proc_macro_attribute]
 pub fn pyclass(attr: TokenStream, input: TokenStream) -> TokenStream {
-    use syn::Item;
     let item = parse_macro_input!(input as Item);
     match item {
         Item::Struct(struct_) => pyclass_impl(attr, struct_, methods_type()),