Skip to content

Commit 268a39f

Browse files
youknowoneyouknowone
authored
Merge pull request RustPython#4901 from youknowone/py-module-def
PyModuleDef and #[pymodule(with(...))]
2 parents cbcdcf4 + 317f432 commit 268a39f

File tree

16 files changed

+316
-115
lines changed

16 files changed

+316
-115
lines changed

‎derive-impl/src/pymodule.rs

Lines changed: 150 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use crate::error::Diagnostic;
22
use crate::util::{
33
format_doc, iter_use_idents, pyclass_ident_and_attrs, text_signature, AttrItemMeta,
44
AttributeExt, ClassItemMeta, ContentItem, ContentItemInner, ErrorVec, ItemMeta, ItemNursery,
5-
SimpleItemMeta, ALL_ALLOWED_NAMES,
5+
ModuleItemMeta, SimpleItemMeta, ALL_ALLOWED_NAMES,
66
};
77
use proc_macro2::{Span, TokenStream};
88
use quote::{quote, quote_spanned, ToTokens};
@@ -45,7 +45,9 @@ impl FromStr for AttrName {
4545
#[derive(Default)]
4646
struct ModuleContext {
4747
name: String,
48-
module_extend_items: ItemNursery,
48+
function_items: FunctionNursery,
49+
attribute_items: ItemNursery,
50+
has_extend_module: bool, // TODO: check if `fn extend_module` exists
4951
errors: Vec<syn::Error>,
5052
}
5153

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

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

9294
// append additional items
9395
let module_name = context.name.as_str();
94-
let module_extend_items = context.module_extend_items.validate()?;
96+
let function_items = context.function_items.validate()?;
97+
let attribute_items = context.attribute_items.validate()?;
9598
let doc = doc.or_else(|| {
9699
crate::doc::Database::shared()
97100
.try_path(module_name)
@@ -104,29 +107,99 @@ pub fn impl_pymodule(attr: AttributeArgs, module_item: Item) -> Result<TokenStre
104107
} else {
105108
quote!(None)
106109
};
110+
let is_submodule = module_meta.sub()?;
111+
let withs = module_meta.with()?;
112+
if !is_submodule {
113+
items.extend(iter_chain![
114+
parse_quote! {
115+
pub(crate) const MODULE_NAME: &'static str = #module_name;
116+
},
117+
parse_quote! {
118+
pub(crate) const DOC: Option<&'static str> = #doc;
119+
},
120+
parse_quote! {
121+
pub(crate) fn __module_def(
122+
ctx: &::rustpython_vm::Context,
123+
) -> &'static ::rustpython_vm::builtins::PyModuleDef {
124+
DEF.get_or_init(|| {
125+
let mut def = ::rustpython_vm::builtins::PyModuleDef {
126+
name: ctx.intern_str(MODULE_NAME),
127+
doc: DOC.map(|doc| ctx.intern_str(doc)),
128+
slots: Default::default(),
129+
};
130+
def.slots.exec = Some(extend_module);
131+
def
132+
})
133+
}
134+
},
135+
parse_quote! {
136+
#[allow(dead_code)]
137+
pub(crate) fn make_module(
138+
vm: &::rustpython_vm::VirtualMachine
139+
) -> ::rustpython_vm::PyRef<::rustpython_vm::builtins::PyModule> {
140+
use ::rustpython_vm::PyPayload;
141+
let module = ::rustpython_vm::builtins::PyModule::from_def(__module_def(&vm.ctx)).into_ref(&vm.ctx);
142+
__init_dict(vm, &module);
143+
extend_module(vm, &module).unwrap();
144+
module
145+
}
146+
},
147+
]);
148+
}
149+
if !is_submodule && !context.has_extend_module {
150+
items.push(parse_quote! {
151+
pub(crate) fn extend_module(vm: &::rustpython_vm::VirtualMachine, module: &::rustpython_vm::Py<::rustpython_vm::builtins::PyModule>) -> ::rustpython_vm::PyResult<()> {
152+
__extend_module(vm, module);
153+
Ok(())
154+
}
155+
});
156+
}
107157
items.extend(iter_chain![
108158
parse_quote! {
109-
pub(crate) const MODULE_NAME: &'static str = #module_name;
159+
::rustpython_vm::common::static_cell! {
160+
pub(crate) static DEF: ::rustpython_vm::builtins::PyModuleDef;
161+
}
162+
},
163+
parse_quote! {
164+
pub(crate) fn __init_attributes(
165+
vm: &::rustpython_vm::VirtualMachine,
166+
module: &::rustpython_vm::Py<::rustpython_vm::builtins::PyModule>,
167+
) {
168+
#(
169+
super::#withs::__init_attributes(vm, module);
170+
)*
171+
let ctx = &vm.ctx;
172+
#attribute_items
173+
}
110174
},
111175
parse_quote! {
112-
pub(crate) const DOC: Option<&'static str> = #doc;
176+
pub(crate) fn __extend_module(
177+
vm: &::rustpython_vm::VirtualMachine,
178+
module: &::rustpython_vm::Py<::rustpython_vm::builtins::PyModule>,
179+
) {
180+
__init_methods(vm, module);
181+
__init_attributes(vm, module);
182+
}
113183
},
114184
parse_quote! {
115-
pub(crate) fn extend_module(
185+
// TODO: remove once PyMethodDef done
186+
pub(crate) fn __init_methods(
116187
vm: &::rustpython_vm::VirtualMachine,
117188
module: &::rustpython_vm::Py<::rustpython_vm::builtins::PyModule>,
118189
) {
119-
#module_extend_items
190+
#(
191+
super::#withs::__init_methods(vm, module);
192+
)*
193+
let ctx = &vm.ctx;
194+
#function_items
120195
}
121196
},
122197
parse_quote! {
123-
#[allow(dead_code)]
124-
pub(crate) fn make_module(
125-
vm: &::rustpython_vm::VirtualMachine
126-
) -> ::rustpython_vm::PyRef<::rustpython_vm::builtins::PyModule> {
127-
let module = vm.new_module(MODULE_NAME, vm.ctx.new_dict(), DOC);
128-
extend_module(vm, &module);
129-
module
198+
pub(crate) fn __init_dict(
199+
vm: &::rustpython_vm::VirtualMachine,
200+
module: &::rustpython_vm::Py<::rustpython_vm::builtins::PyModule>,
201+
) {
202+
::rustpython_vm::builtins::PyModule::__init_dict_from_def(vm, module);
130203
}
131204
},
132205
]);
@@ -248,6 +321,53 @@ where
248321
Ok((result, cfgs))
249322
}
250323

324+
#[derive(Default)]
325+
struct FunctionNursery {
326+
items: Vec<FunctionNurseryItem>,
327+
}
328+
329+
struct FunctionNurseryItem {
330+
py_names: Vec<String>,
331+
cfgs: Vec<Attribute>,
332+
ident: Ident,
333+
#[allow(dead_code)]
334+
doc: String,
335+
tokens: TokenStream,
336+
}
337+
338+
impl FunctionNursery {
339+
fn add_item(&mut self, item: FunctionNurseryItem) {
340+
self.items.push(item);
341+
}
342+
343+
fn validate(self) -> Result<ValidatedFunctionNursery> {
344+
let mut name_set = HashSet::new();
345+
for item in &self.items {
346+
for py_name in &item.py_names {
347+
if !name_set.insert((py_name.to_owned(), &item.cfgs)) {
348+
bail_span!(item.ident, "duplicate method name `{}`", py_name);
349+
}
350+
}
351+
}
352+
Ok(ValidatedFunctionNursery(self))
353+
}
354+
}
355+
356+
struct ValidatedFunctionNursery(FunctionNursery);
357+
358+
impl ToTokens for ValidatedFunctionNursery {
359+
fn to_tokens(&self, tokens: &mut TokenStream) {
360+
for item in &self.0.items {
361+
let cfgs = &item.cfgs;
362+
let item_tokens = &item.tokens;
363+
tokens.extend(quote! {
364+
#(#cfgs)*
365+
#item_tokens
366+
});
367+
}
368+
}
369+
}
370+
251371
/// #[pyfunction]
252372
struct FunctionItem {
253373
inner: ContentItemInner<AttrName>,
@@ -318,7 +438,7 @@ impl ModuleItem for FunctionItem {
318438
let py_name = item_meta.simple_name()?;
319439
let sig_doc = text_signature(func.sig(), &py_name);
320440

321-
let (tokens, py_names) = {
441+
let (tokens, py_names, doc) = {
322442
let module = args.module_name();
323443
let doc = args.attrs.doc().or_else(|| {
324444
crate::doc::Database::shared()
@@ -332,10 +452,10 @@ impl ModuleItem for FunctionItem {
332452
} else {
333453
sig_doc
334454
};
335-
let doc = quote!(.with_doc(#doc.to_owned(), &vm.ctx));
455+
let with_doc = quote!(.with_doc(#doc.to_owned(), &vm.ctx));
336456
let new_func = quote_spanned!(ident.span()=>
337457
vm.ctx.make_func_def(vm.ctx.intern_str(#py_name), #ident)
338-
#doc
458+
#with_doc
339459
.into_function()
340460
.with_module(vm.new_pyobj(#module.to_owned()))
341461
.into_ref(&vm.ctx)
@@ -348,6 +468,7 @@ impl ModuleItem for FunctionItem {
348468
vm.__module_set_attr(module, #py_name, func).unwrap();
349469
}},
350470
vec![py_name],
471+
doc,
351472
)
352473
} else {
353474
let mut py_names = HashSet::new();
@@ -381,17 +502,18 @@ impl ModuleItem for FunctionItem {
381502
}
382503
}},
383504
py_names,
505+
doc,
384506
)
385507
}
386508
};
387509

388-
args.context.module_extend_items.add_item(
389-
ident.clone(),
510+
args.context.function_items.add_item(FunctionNurseryItem {
511+
ident: ident.to_owned(),
390512
py_names,
391-
args.cfgs.to_vec(),
513+
cfgs: args.cfgs.to_vec(),
514+
doc,
392515
tokens,
393-
10,
394-
)?;
516+
});
395517
Ok(())
396518
}
397519
}
@@ -432,8 +554,8 @@ impl ModuleItem for ClassItem {
432554
class_meta.class_name()?
433555
};
434556
let class_new = quote_spanned!(ident.span() =>
435-
let new_class = <#ident as ::rustpython_vm::class::PyClassImpl>::make_class(&vm.ctx);
436-
new_class.set_attr(rustpython_vm::identifier!(vm, __module__), vm.new_pyobj(#module_name));
557+
let new_class = <#ident as ::rustpython_vm::class::PyClassImpl>::make_class(ctx);
558+
new_class.set_attr(rustpython_vm::identifier!(ctx, __module__), vm.new_pyobj(#module_name));
437559
);
438560
(class_name, class_new)
439561
};
@@ -473,7 +595,7 @@ impl ModuleItem for ClassItem {
473595
},
474596
};
475597

476-
args.context.module_extend_items.add_item(
598+
args.context.attribute_items.add_item(
477599
ident.clone(),
478600
py_names,
479601
args.cfgs.to_vec(),
@@ -561,7 +683,7 @@ impl ModuleItem for AttributeItem {
561683
let tokens = quote_spanned! { ident.span() =>
562684
vm.__module_set_attr(module, #py_name, vm.new_pyobj(#ident)).unwrap();
563685
};
564-
args.context.module_extend_items.add_item(
686+
args.context.attribute_items.add_item(
565687
ident.clone(),
566688
vec![py_name],
567689
cfgs.clone(),
@@ -624,7 +746,7 @@ impl ModuleItem for AttributeItem {
624746
};
625747

626748
args.context
627-
.module_extend_items
749+
.attribute_items
628750
.add_item(ident, py_names, cfgs, tokens, 1)?;
629751

630752
Ok(())

‎derive-impl/src/util.rs

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,21 @@ impl ItemMetaInner {
197197
};
198198
Ok(value)
199199
}
200+
201+
pub fn _optional_list(
202+
&self,
203+
key: &str,
204+
) -> Result<Option<impl std::iter::Iterator<Item = &'_ NestedMeta>>> {
205+
let value = if let Some((_, meta)) = self.meta_map.get(key) {
206+
let Meta::List(syn::MetaList { path: _, nested, .. }) = meta else {
207+
bail_span!(meta, "#[{}({}(...))] must be a list", self.meta_name(), key)
208+
};
209+
Some(nested.into_iter())
210+
} else {
211+
None
212+
};
213+
Ok(value)
214+
}
200215
}
201216

202217
pub(crate) trait ItemMeta: Sized {
@@ -251,6 +266,38 @@ impl ItemMeta for SimpleItemMeta {
251266
}
252267
}
253268

269+
pub(crate) struct ModuleItemMeta(pub ItemMetaInner);
270+
271+
impl ItemMeta for ModuleItemMeta {
272+
const ALLOWED_NAMES: &'static [&'static str] = &["name", "with", "sub"];
273+
274+
fn from_inner(inner: ItemMetaInner) -> Self {
275+
Self(inner)
276+
}
277+
fn inner(&self) -> &ItemMetaInner {
278+
&self.0
279+
}
280+
}
281+
282+
impl ModuleItemMeta {
283+
pub fn sub(&self) -> Result<bool> {
284+
self.inner()._bool("sub")
285+
}
286+
pub fn with(&self) -> Result<Vec<&syn::Path>> {
287+
let mut withs = Vec::new();
288+
let Some(nested) = self.inner()._optional_list("with")? else {
289+
return Ok(withs);
290+
};
291+
for meta in nested {
292+
let NestedMeta::Meta(Meta::Path(path)) = meta else {
293+
bail_span!(meta, "#[pymodule(with(...))] arguments should be paths")
294+
};
295+
withs.push(path);
296+
}
297+
Ok(withs)
298+
}
299+
}
300+
254301
pub(crate) struct AttrItemMeta(pub ItemMetaInner);
255302

256303
impl ItemMeta for AttrItemMeta {

0 commit comments

Comments
 (0)