Skip to content

PyModuleDef and #[pymodule(with(...))] #4901

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Apr 24, 2023
Merged

PyModuleDef and #[pymodule(with(...))] #4901

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
178 changes: 150 additions & 28 deletions derive-impl/src/pymodule.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use crate::error::Diagnostic;
use crate::util::{
format_doc, iter_use_idents, pyclass_ident_and_attrs, text_signature, AttrItemMeta,
AttributeExt, ClassItemMeta, ContentItem, ContentItemInner, ErrorVec, ItemMeta, ItemNursery,
SimpleItemMeta, ALL_ALLOWED_NAMES,
ModuleItemMeta, SimpleItemMeta, ALL_ALLOWED_NAMES,
};
use proc_macro2::{Span, TokenStream};
use quote::{quote, quote_spanned, ToTokens};
Expand Down Expand Up @@ -45,7 +45,9 @@ impl FromStr for AttrName {
#[derive(Default)]
struct ModuleContext {
name: String,
module_extend_items: ItemNursery,
function_items: FunctionNursery,
attribute_items: ItemNursery,
has_extend_module: bool, // TODO: check if `fn extend_module` exists
errors: Vec<syn::Error>,
}

Expand All @@ -56,7 +58,7 @@ pub fn impl_pymodule(attr: AttributeArgs, module_item: Item) -> Result<TokenStre
};
let fake_ident = Ident::new("pymodule", module_item.span());
let module_meta =
SimpleItemMeta::from_nested(module_item.ident.clone(), fake_ident, attr.into_iter())?;
ModuleItemMeta::from_nested(module_item.ident.clone(), fake_ident, attr.into_iter())?;

// generation resources
let mut context = ModuleContext {
Expand Down Expand Up @@ -91,7 +93,8 @@ pub fn impl_pymodule(attr: AttributeArgs, module_item: Item) -> Result<TokenStre

// append additional items
let module_name = context.name.as_str();
let module_extend_items = context.module_extend_items.validate()?;
let function_items = context.function_items.validate()?;
let attribute_items = context.attribute_items.validate()?;
let doc = doc.or_else(|| {
crate::doc::Database::shared()
.try_path(module_name)
Expand All @@ -104,29 +107,99 @@ pub fn impl_pymodule(attr: AttributeArgs, module_item: Item) -> Result<TokenStre
} else {
quote!(None)
};
let is_submodule = module_meta.sub()?;
let withs = module_meta.with()?;
if !is_submodule {
items.extend(iter_chain![
parse_quote! {
pub(crate) const MODULE_NAME: &'static str = #module_name;
},
parse_quote! {
pub(crate) const DOC: Option<&'static str> = #doc;
},
parse_quote! {
pub(crate) fn __module_def(
ctx: &::rustpython_vm::Context,
) -> &'static ::rustpython_vm::builtins::PyModuleDef {
DEF.get_or_init(|| {
let mut def = ::rustpython_vm::builtins::PyModuleDef {
name: ctx.intern_str(MODULE_NAME),
doc: DOC.map(|doc| ctx.intern_str(doc)),
slots: Default::default(),
};
def.slots.exec = Some(extend_module);
def
})
}
},
parse_quote! {
#[allow(dead_code)]
pub(crate) fn make_module(
vm: &::rustpython_vm::VirtualMachine
) -> ::rustpython_vm::PyRef<::rustpython_vm::builtins::PyModule> {
use ::rustpython_vm::PyPayload;
let module = ::rustpython_vm::builtins::PyModule::from_def(__module_def(&vm.ctx)).into_ref(&vm.ctx);
__init_dict(vm, &module);
extend_module(vm, &module).unwrap();
module
}
},
]);
}
if !is_submodule && !context.has_extend_module {
items.push(parse_quote! {
pub(crate) fn extend_module(vm: &::rustpython_vm::VirtualMachine, module: &::rustpython_vm::Py<::rustpython_vm::builtins::PyModule>) -> ::rustpython_vm::PyResult<()> {
__extend_module(vm, module);
Ok(())
}
});
}
items.extend(iter_chain![
parse_quote! {
pub(crate) const MODULE_NAME: &'static str = #module_name;
::rustpython_vm::common::static_cell! {
pub(crate) static DEF: ::rustpython_vm::builtins::PyModuleDef;
}
},
parse_quote! {
pub(crate) fn __init_attributes(
vm: &::rustpython_vm::VirtualMachine,
module: &::rustpython_vm::Py<::rustpython_vm::builtins::PyModule>,
) {
#(
super::#withs::__init_attributes(vm, module);
)*
let ctx = &vm.ctx;
#attribute_items
}
},
parse_quote! {
pub(crate) const DOC: Option<&'static str> = #doc;
pub(crate) fn __extend_module(
vm: &::rustpython_vm::VirtualMachine,
module: &::rustpython_vm::Py<::rustpython_vm::builtins::PyModule>,
) {
__init_methods(vm, module);
__init_attributes(vm, module);
}
},
parse_quote! {
pub(crate) fn extend_module(
// TODO: remove once PyMethodDef done
pub(crate) fn __init_methods(
vm: &::rustpython_vm::VirtualMachine,
module: &::rustpython_vm::Py<::rustpython_vm::builtins::PyModule>,
) {
#module_extend_items
#(
super::#withs::__init_methods(vm, module);
)*
let ctx = &vm.ctx;
#function_items
}
},
parse_quote! {
#[allow(dead_code)]
pub(crate) fn make_module(
vm: &::rustpython_vm::VirtualMachine
) -> ::rustpython_vm::PyRef<::rustpython_vm::builtins::PyModule> {
let module = vm.new_module(MODULE_NAME, vm.ctx.new_dict(), DOC);
extend_module(vm, &module);
module
pub(crate) fn __init_dict(
vm: &::rustpython_vm::VirtualMachine,
module: &::rustpython_vm::Py<::rustpython_vm::builtins::PyModule>,
) {
::rustpython_vm::builtins::PyModule::__init_dict_from_def(vm, module);
}
},
]);
Expand Down Expand Up @@ -248,6 +321,53 @@ where
Ok((result, cfgs))
}

#[derive(Default)]
struct FunctionNursery {
items: Vec<FunctionNurseryItem>,
}

struct FunctionNurseryItem {
py_names: Vec<String>,
cfgs: Vec<Attribute>,
ident: Ident,
#[allow(dead_code)]
doc: String,
tokens: TokenStream,
}

impl FunctionNursery {
fn add_item(&mut self, item: FunctionNurseryItem) {
self.items.push(item);
}

fn validate(self) -> Result<ValidatedFunctionNursery> {
let mut name_set = HashSet::new();
for item in &self.items {
for py_name in &item.py_names {
if !name_set.insert((py_name.to_owned(), &item.cfgs)) {
bail_span!(item.ident, "duplicate method name `{}`", py_name);
}
}
}
Ok(ValidatedFunctionNursery(self))
}
}

struct ValidatedFunctionNursery(FunctionNursery);

impl ToTokens for ValidatedFunctionNursery {
fn to_tokens(&self, tokens: &mut TokenStream) {
for item in &self.0.items {
let cfgs = &item.cfgs;
let item_tokens = &item.tokens;
tokens.extend(quote! {
#(#cfgs)*
#item_tokens
});
}
}
}

/// #[pyfunction]
struct FunctionItem {
inner: ContentItemInner<AttrName>,
Expand Down Expand Up @@ -318,7 +438,7 @@ impl ModuleItem for FunctionItem {
let py_name = item_meta.simple_name()?;
let sig_doc = text_signature(func.sig(), &py_name);

let (tokens, py_names) = {
let (tokens, py_names, doc) = {
let module = args.module_name();
let doc = args.attrs.doc().or_else(|| {
crate::doc::Database::shared()
Expand All @@ -332,10 +452,10 @@ impl ModuleItem for FunctionItem {
} else {
sig_doc
};
let doc = quote!(.with_doc(#doc.to_owned(), &vm.ctx));
let with_doc = quote!(.with_doc(#doc.to_owned(), &vm.ctx));
let new_func = quote_spanned!(ident.span()=>
vm.ctx.make_func_def(vm.ctx.intern_str(#py_name), #ident)
#doc
#with_doc
.into_function()
.with_module(vm.new_pyobj(#module.to_owned()))
.into_ref(&vm.ctx)
Expand All @@ -348,6 +468,7 @@ impl ModuleItem for FunctionItem {
vm.__module_set_attr(module, #py_name, func).unwrap();
}},
vec![py_name],
doc,
)
} else {
let mut py_names = HashSet::new();
Expand Down Expand Up @@ -381,17 +502,18 @@ impl ModuleItem for FunctionItem {
}
}},
py_names,
doc,
)
}
};

args.context.module_extend_items.add_item(
ident.clone(),
args.context.function_items.add_item(FunctionNurseryItem {
ident: ident.to_owned(),
py_names,
args.cfgs.to_vec(),
cfgs: args.cfgs.to_vec(),
doc,
tokens,
10,
)?;
});
Ok(())
}
}
Expand Down Expand Up @@ -432,8 +554,8 @@ impl ModuleItem for ClassItem {
class_meta.class_name()?
};
let class_new = quote_spanned!(ident.span() =>
let new_class = <#ident as ::rustpython_vm::class::PyClassImpl>::make_class(&vm.ctx);
new_class.set_attr(rustpython_vm::identifier!(vm, __module__), vm.new_pyobj(#module_name));
let new_class = <#ident as ::rustpython_vm::class::PyClassImpl>::make_class(ctx);
new_class.set_attr(rustpython_vm::identifier!(ctx, __module__), vm.new_pyobj(#module_name));
);
(class_name, class_new)
};
Expand Down Expand Up @@ -473,7 +595,7 @@ impl ModuleItem for ClassItem {
},
};

args.context.module_extend_items.add_item(
args.context.attribute_items.add_item(
ident.clone(),
py_names,
args.cfgs.to_vec(),
Expand Down Expand Up @@ -561,7 +683,7 @@ impl ModuleItem for AttributeItem {
let tokens = quote_spanned! { ident.span() =>
vm.__module_set_attr(module, #py_name, vm.new_pyobj(#ident)).unwrap();
};
args.context.module_extend_items.add_item(
args.context.attribute_items.add_item(
ident.clone(),
vec![py_name],
cfgs.clone(),
Expand Down Expand Up @@ -624,7 +746,7 @@ impl ModuleItem for AttributeItem {
};

args.context
.module_extend_items
.attribute_items
.add_item(ident, py_names, cfgs, tokens, 1)?;

Ok(())
Expand Down
47 changes: 47 additions & 0 deletions derive-impl/src/util.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,21 @@ impl ItemMetaInner {
};
Ok(value)
}

pub fn _optional_list(
&self,
key: &str,
) -> Result<Option<impl std::iter::Iterator<Item = &'_ NestedMeta>>> {
let value = if let Some((_, meta)) = self.meta_map.get(key) {
let Meta::List(syn::MetaList { path: _, nested, .. }) = meta else {
bail_span!(meta, "#[{}({}(...))] must be a list", self.meta_name(), key)
};
Some(nested.into_iter())
} else {
None
};
Ok(value)
}
}

pub(crate) trait ItemMeta: Sized {
Expand Down Expand Up @@ -251,6 +266,38 @@ impl ItemMeta for SimpleItemMeta {
}
}

pub(crate) struct ModuleItemMeta(pub ItemMetaInner);

impl ItemMeta for ModuleItemMeta {
const ALLOWED_NAMES: &'static [&'static str] = &["name", "with", "sub"];

fn from_inner(inner: ItemMetaInner) -> Self {
Self(inner)
}
fn inner(&self) -> &ItemMetaInner {
&self.0
}
}

impl ModuleItemMeta {
pub fn sub(&self) -> Result<bool> {
self.inner()._bool("sub")
}
pub fn with(&self) -> Result<Vec<&syn::Path>> {
let mut withs = Vec::new();
let Some(nested) = self.inner()._optional_list("with")? else {
return Ok(withs);
};
for meta in nested {
let NestedMeta::Meta(Meta::Path(path)) = meta else {
bail_span!(meta, "#[pymodule(with(...))] arguments should be paths")
};
withs.push(path);
}
Ok(withs)
}
}

pub(crate) struct AttrItemMeta(pub ItemMetaInner);

impl ItemMeta for AttrItemMeta {
Expand Down
Loading