Skip to content

Commit 98eef08

Browse files
youknowoneyouknowone
committed
Method overhaul with PyMethodDef
1 parent 4a099ab commit 98eef08

35 files changed

+1128
-679
lines changed

‎derive-impl/src/pyclass.rs‎

Lines changed: 129 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use crate::util::{
66
};
77
use proc_macro2::{Span, TokenStream};
88
use quote::{quote, quote_spanned, ToTokens};
9-
use std::collections::HashMap;
9+
use std::collections::{HashMap, HashSet};
1010
use std::str::FromStr;
1111
use syn::{
1212
parse_quote, spanned::Spanned, Attribute, AttributeArgs, Ident, Item, Meta, NestedMeta, Result,
@@ -62,7 +62,8 @@ impl FromStr for AttrName {
6262

6363
#[derive(Default)]
6464
struct ImplContext {
65-
impl_extend_items: ItemNursery,
65+
attribute_items: ItemNursery,
66+
method_items: MethodNursery,
6667
getset_items: GetSetNursery,
6768
member_items: MemberNursery,
6869
extend_slots_items: ItemNursery,
@@ -92,6 +93,7 @@ fn extract_items_into_context<'a, Item>(
9293
});
9394
context.errors.ok_or_push(r);
9495
}
96+
context.errors.ok_or_push(context.method_items.validate());
9597
context.errors.ok_or_push(context.getset_items.validate());
9698
context.errors.ok_or_push(context.member_items.validate());
9799
}
@@ -157,18 +159,28 @@ pub(crate) fn impl_pyclass_impl(attr: AttributeArgs, item: Item) -> Result<Token
157159

158160
let ExtractedImplAttrs {
159161
payload: attr_payload,
160-
with_impl,
161162
flags,
163+
with_impl,
164+
with_method_defs,
162165
with_slots,
163166
} = extract_impl_attrs(attr, &impl_ty)?;
164167
let payload_ty = attr_payload.unwrap_or(payload_guess);
168+
let method_def = &context.method_items;
165169
let getset_impl = &context.getset_items;
166170
let member_impl = &context.member_items;
167-
let extend_impl = context.impl_extend_items.validate()?;
171+
let extend_impl = context.attribute_items.validate()?;
168172
let slots_impl = context.extend_slots_items.validate()?;
169173
let class_extensions = &context.class_extensions;
170174

171175
let extra_methods = iter_chain![
176+
parse_quote! {
177+
#[allow(clippy::ptr_arg)]
178+
fn __extend_method_def(
179+
method_defs: &mut Vec<::rustpython_vm::function::PyMethodDef>,
180+
) {
181+
#method_def
182+
}
183+
},
172184
parse_quote! {
173185
fn __extend_py_class(
174186
ctx: &::rustpython_vm::Context,
@@ -202,6 +214,14 @@ pub(crate) fn impl_pyclass_impl(attr: AttributeArgs, item: Item) -> Result<Token
202214
#with_impl
203215
}
204216

217+
#[allow(clippy::ptr_arg)]
218+
fn impl_extend_method_def(
219+
method_defs: &mut Vec<::rustpython_vm::function::PyMethodDef>,
220+
) {
221+
#impl_ty::__extend_method_def(method_defs);
222+
#with_method_defs
223+
}
224+
205225
fn extend_slots(slots: &mut ::rustpython_vm::types::PyTypeSlots) {
206226
#impl_ty::__extend_slots(slots);
207227
#with_slots
@@ -235,9 +255,10 @@ pub(crate) fn impl_pyclass_impl(attr: AttributeArgs, item: Item) -> Result<Token
235255
..
236256
} = extract_impl_attrs(attr, &trai.ident)?;
237257

258+
let method_def = &context.method_items;
238259
let getset_impl = &context.getset_items;
239260
let member_impl = &context.member_items;
240-
let extend_impl = &context.impl_extend_items.validate()?;
261+
let extend_impl = &context.attribute_items.validate()?;
241262
let slots_impl = &context.extend_slots_items.validate()?;
242263
let class_extensions = &context.class_extensions;
243264
let call_extend_slots = if has_extend_slots {
@@ -248,6 +269,14 @@ pub(crate) fn impl_pyclass_impl(attr: AttributeArgs, item: Item) -> Result<Token
248269
quote! {}
249270
};
250271
let extra_methods = iter_chain![
272+
parse_quote! {
273+
#[allow(clippy::ptr_arg)]
274+
fn __extend_method_def(
275+
method_defs: &mut Vec<::rustpython_vm::function::PyMethodDef>,
276+
) {
277+
#method_def
278+
}
279+
},
251280
parse_quote! {
252281
fn __extend_py_class(
253282
ctx: &::rustpython_vm::Context,
@@ -732,51 +761,14 @@ where
732761
let py_name = item_meta.method_name()?;
733762
let sig_doc = text_signature(func.sig(), &py_name);
734763

735-
let tokens = {
736-
let doc = args.attrs.doc().map_or_else(TokenStream::new, |mut doc| {
737-
doc = format_doc(&sig_doc, &doc);
738-
quote!(.with_doc(#doc.to_owned(), ctx))
739-
});
740-
let build_func = match self.inner.attr_name {
741-
AttrName::Method => quote!(.build_method(ctx, class, false)),
742-
AttrName::ClassMethod => quote!(.build_classmethod(ctx, class)),
743-
AttrName::StaticMethod => quote!(.build_staticmethod(ctx, class)),
744-
other => unreachable!(
745-
"Only 'method', 'classmethod' and 'staticmethod' are supported, got {:?}",
746-
other
747-
),
748-
};
749-
if py_name.starts_with("__") && py_name.ends_with("__") {
750-
let name_ident = Ident::new(&py_name, ident.span());
751-
quote_spanned! { ident.span() =>
752-
class.set_attr(
753-
ctx.names.#name_ident,
754-
ctx.make_func_def(ctx.intern_str(#py_name), Self::#ident)
755-
#doc
756-
#build_func
757-
.into(),
758-
);
759-
}
760-
} else {
761-
quote_spanned! { ident.span() =>
762-
class.set_str_attr(
763-
#py_name,
764-
ctx.make_func_def(ctx.intern_str(#py_name), Self::#ident)
765-
#doc
766-
#build_func,
767-
ctx,
768-
);
769-
}
770-
}
771-
};
772-
773-
args.context.impl_extend_items.add_item(
774-
ident.clone(),
775-
vec![py_name],
776-
args.cfgs.to_vec(),
777-
tokens,
778-
5,
779-
)?;
764+
let doc = args.attrs.doc().map(|doc| format_doc(&sig_doc, &doc));
765+
args.context.method_items.add_item(MethodNurseryItem {
766+
py_name,
767+
cfgs: args.cfgs.to_vec(),
768+
ident: ident.to_owned(),
769+
doc,
770+
attr_name: self.inner.attr_name,
771+
});
780772
Ok(())
781773
}
782774
}
@@ -898,7 +890,7 @@ where
898890
};
899891

900892
args.context
901-
.impl_extend_items
893+
.attribute_items
902894
.add_item(ident.clone(), vec![py_name], cfgs, tokens, 1)?;
903895

904896
Ok(())
@@ -960,6 +952,78 @@ where
960952
}
961953
}
962954

955+
#[derive(Default)]
956+
struct MethodNursery {
957+
items: Vec<MethodNurseryItem>,
958+
}
959+
960+
struct MethodNurseryItem {
961+
py_name: String,
962+
cfgs: Vec<Attribute>,
963+
ident: Ident,
964+
doc: Option<String>,
965+
attr_name: AttrName,
966+
}
967+
968+
impl MethodNursery {
969+
fn add_item(&mut self, item: MethodNurseryItem) {
970+
self.items.push(item);
971+
}
972+
973+
fn validate(&mut self) -> Result<()> {
974+
let mut name_set = HashSet::new();
975+
for item in &self.items {
976+
if !name_set.insert((&item.py_name, &item.cfgs)) {
977+
bail_span!(item.ident, "duplicate method name `{}`", item.py_name);
978+
}
979+
}
980+
Ok(())
981+
}
982+
}
983+
984+
impl ToTokens for MethodNursery {
985+
fn to_tokens(&self, tokens: &mut TokenStream) {
986+
for item in &self.items {
987+
let py_name = &item.py_name;
988+
let ident = &item.ident;
989+
let cfgs = &item.cfgs;
990+
let doc = if let Some(doc) = item.doc.as_ref() {
991+
quote! { Some(#doc) }
992+
} else {
993+
quote! { None }
994+
};
995+
let flags = match &item.attr_name {
996+
AttrName::Method => {
997+
quote! { rustpython_vm::function::PyMethodFlags::METHOD }
998+
}
999+
AttrName::ClassMethod => {
1000+
quote! { rustpython_vm::function::PyMethodFlags::CLASS }
1001+
}
1002+
AttrName::StaticMethod => {
1003+
quote! { rustpython_vm::function::PyMethodFlags::STATIC }
1004+
}
1005+
_ => unreachable!(),
1006+
};
1007+
// TODO: intern
1008+
// let py_name = if py_name.starts_with("__") && py_name.ends_with("__") {
1009+
// let name_ident = Ident::new(&py_name, ident.span());
1010+
// quote_spanned! { ident.span() => ctx.names.#name_ident }
1011+
// } else {
1012+
// quote_spanned! { ident.span() => #py_name }
1013+
// };
1014+
tokens.extend(quote! {
1015+
#(#cfgs)*
1016+
method_defs.push(rustpython_vm::function::PyMethodDef {
1017+
name: #py_name,
1018+
func: rustpython_vm::function::IntoPyNativeFn::into_func(Self::#ident),
1019+
flags: #flags,
1020+
doc: #doc,
1021+
});
1022+
});
1023+
}
1024+
}
1025+
}
1026+
9631027
#[derive(Default)]
9641028
#[allow(clippy::type_complexity)]
9651029
struct GetSetNursery {
@@ -1367,13 +1431,15 @@ impl MemberItemMeta {
13671431

13681432
struct ExtractedImplAttrs {
13691433
payload: Option<Ident>,
1434+
flags: TokenStream,
13701435
with_impl: TokenStream,
1436+
with_method_defs: TokenStream,
13711437
with_slots: TokenStream,
1372-
flags: TokenStream,
13731438
}
13741439

13751440
fn extract_impl_attrs(attr: AttributeArgs, item: &Ident) -> Result<ExtractedImplAttrs> {
13761441
let mut withs = Vec::new();
1442+
let mut with_method_defs = Vec::new();
13771443
let mut with_slots = Vec::new();
13781444
let mut flags = vec![quote! {
13791445
{
@@ -1396,23 +1462,28 @@ fn extract_impl_attrs(attr: AttributeArgs, item: &Ident) -> Result<ExtractedImpl
13961462
let NestedMeta::Meta(Meta::Path(path)) = meta else {
13971463
bail_span!(meta, "#[pyclass(with(...))] arguments should be paths")
13981464
};
1399-
let (extend_class, extend_slots) =
1465+
let (extend_class, extend_method_def, extend_slots) =
14001466
if path.is_ident("PyRef") || path.is_ident("Py") {
14011467
// special handling for PyRef
14021468
(
14031469
quote!(#path::<Self>::__extend_py_class),
1470+
quote!(#path::<Self>::__extend_method_def),
14041471
quote!(#path::<Self>::__extend_slots),
14051472
)
14061473
} else {
14071474
(
14081475
quote!(<Self as #path>::__extend_py_class),
1476+
quote!(<Self as #path>::__extend_method_def),
14091477
quote!(<Self as #path>::__extend_slots),
14101478
)
14111479
};
14121480
let item_span = item.span().resolved_at(Span::call_site());
14131481
withs.push(quote_spanned! { path.span() =>
14141482
#extend_class(ctx, class);
14151483
});
1484+
with_method_defs.push(quote_spanned! { path.span() =>
1485+
#extend_method_def(method_defs);
1486+
});
14161487
with_slots.push(quote_spanned! { item_span =>
14171488
#extend_slots(slots);
14181489
});
@@ -1450,11 +1521,14 @@ fn extract_impl_attrs(attr: AttributeArgs, item: &Ident) -> Result<ExtractedImpl
14501521

14511522
Ok(ExtractedImplAttrs {
14521523
payload,
1524+
flags: quote! {
1525+
#(#flags)*
1526+
},
14531527
with_impl: quote! {
14541528
#(#withs)*
14551529
},
1456-
flags: quote! {
1457-
#(#flags)*
1530+
with_method_defs: quote! {
1531+
#(#with_method_defs)*
14581532
},
14591533
with_slots: quote! {
14601534
#(#with_slots)*

0 commit comments

Comments
 (0)