Skip to content

Commit

Permalink
#[pymodule] mod some_module { ... } v3
Browse files Browse the repository at this point in the history
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>
  • Loading branch information
3 people committed Feb 9, 2024
1 parent 9bb0011 commit bd7ad6b
Showing 20 changed files with 403 additions and 56 deletions.
1 change: 1 addition & 0 deletions newsfragments/3815.added.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
The ability to create Python modules with a Rust `mod` block.
4 changes: 3 additions & 1 deletion pyo3-macros-backend/src/lib.rs
Original file line number Diff line number Diff line change
@@ -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};
160 changes: 156 additions & 4 deletions pyo3-macros-backend/src/module.rs
Original file line number Diff line number Diff line change
@@ -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)
}
}
}
4 changes: 4 additions & 0 deletions pyo3-macros-backend/src/pyfunction.rs
Original file line number Diff line number Diff line change
@@ -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 -
12 changes: 9 additions & 3 deletions pyo3-macros/src/lib.rs
Original file line number Diff line number Diff line change
@@ -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
92 changes: 53 additions & 39 deletions pytests/src/lib.rs
Original file line number Diff line number Diff line change
@@ -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(())
}
}
13 changes: 12 additions & 1 deletion src/impl_/pymodule.rs
Original file line number Diff line number Diff line change
@@ -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 {
@@ -132,6 +132,17 @@ impl ModuleDef {
}
}

/// Trait to add an element (class, function...) to a module
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};
4 changes: 2 additions & 2 deletions src/macros.rs
Original file line number Diff line number Diff line change
@@ -167,8 +167,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),
);
}
};
Loading

0 comments on commit bd7ad6b

Please sign in to comment.