Skip to content

Commit 06f8d0b

Browse files
youknowoneyouknowone
committed
pyclass member support for #[pymodule]
1 parent 7cea2c3 commit 06f8d0b

4 files changed

Lines changed: 119 additions & 91 deletions

File tree

‎derive/src/pyclass.rs‎

Lines changed: 36 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
use super::Diagnostic;
2-
use crate::util::{def_to_name, path_eq, strip_prefix, ItemMeta};
2+
use crate::util::{def_to_name, path_eq, strip_prefix, ItemIdent, ItemMeta};
33
use proc_macro2::{Span, TokenStream as TokenStream2};
44
use quote::{quote, quote_spanned, ToTokens};
55
use std::collections::{HashMap, HashSet};
66
use syn::{
77
parse_quote, spanned::Spanned, Attribute, AttributeArgs, Ident, Index, Item, Lit, Meta,
8-
NestedMeta, Signature,
8+
NestedMeta,
99
};
1010

1111
fn meta_to_vec(meta: Meta) -> Result<Vec<NestedMeta>, Meta> {
@@ -54,7 +54,7 @@ impl Class {
5454
}
5555
}
5656

57-
fn extract_method(sig: &Signature, meta: Meta) -> Result<ClassItem, Diagnostic> {
57+
fn extract_method(ident: &Ident, meta: Meta) -> Result<ClassItem, Diagnostic> {
5858
let nesteds = meta_to_vec(meta).map_err(|meta| {
5959
err_span!(
6060
meta,
@@ -64,30 +64,34 @@ impl Class {
6464
})?;
6565

6666
let item_meta =
67-
ItemMeta::from_nested_meta("pymethod", sig, &nesteds, ItemMeta::ATTRIBUTE_NAMES)?;
67+
ItemMeta::from_nested_meta("pymethod", &ident, &nesteds, ItemMeta::ATTRIBUTE_NAMES)?;
6868
Ok(ClassItem::Method {
69-
item_ident: sig.ident.clone(),
69+
item_ident: ident.clone(),
7070
py_name: item_meta.method_name()?,
7171
})
7272
}
7373

74-
fn extract_classmethod(sig: &Signature, meta: Meta) -> Result<ClassItem, Diagnostic> {
74+
fn extract_classmethod(ident: &Ident, meta: Meta) -> Result<ClassItem, Diagnostic> {
7575
let nesteds = meta_to_vec(meta).map_err(|meta| {
7676
err_span!(
7777
meta,
7878
"#[pyclassmethod = \"...\"] cannot be a name/value, you probably meant \
7979
#[pyclassmethod(name = \"...\")]",
8080
)
8181
})?;
82-
let item_meta =
83-
ItemMeta::from_nested_meta("pyclassmethod", sig, &nesteds, ItemMeta::ATTRIBUTE_NAMES)?;
82+
let item_meta = ItemMeta::from_nested_meta(
83+
"pyclassmethod",
84+
&ident,
85+
&nesteds,
86+
ItemMeta::ATTRIBUTE_NAMES,
87+
)?;
8488
Ok(ClassItem::ClassMethod {
85-
item_ident: sig.ident.clone(),
89+
item_ident: ident.clone(),
8690
py_name: item_meta.method_name()?,
8791
})
8892
}
8993

90-
fn extract_property(sig: &Signature, meta: Meta) -> Result<ClassItem, Diagnostic> {
94+
fn extract_property(ident: &Ident, meta: Meta) -> Result<ClassItem, Diagnostic> {
9195
let nesteds = meta_to_vec(meta).map_err(|meta| {
9296
err_span!(
9397
meta,
@@ -96,26 +100,26 @@ impl Class {
96100
)
97101
})?;
98102
let item_meta =
99-
ItemMeta::from_nested_meta("pyproperty", sig, &nesteds, ItemMeta::PROPERTY_NAMES)?;
103+
ItemMeta::from_nested_meta("pyproperty", &ident, &nesteds, ItemMeta::PROPERTY_NAMES)?;
100104
Ok(ClassItem::Property {
101105
py_name: item_meta.property_name()?,
102-
item_ident: sig.ident.clone(),
106+
item_ident: ident.clone(),
103107
setter: item_meta.setter()?,
104108
})
105109
}
106110

107-
fn extract_slot(sig: &Signature, meta: Meta) -> Result<ClassItem, Diagnostic> {
111+
fn extract_slot(ident: &Ident, meta: Meta) -> Result<ClassItem, Diagnostic> {
108112
let pyslot_err = "#[pyslot] must be of the form #[pyslot] or #[pyslot(slotname)]";
109113
let nesteds = meta_to_vec(meta).map_err(|meta| err_span!(meta, "{}", pyslot_err))?;
110114
if nesteds.len() > 1 {
111115
return Err(Diagnostic::spanned_error(&quote!(#(#nesteds)*), pyslot_err));
112116
}
113117
let slot_ident = if nesteds.is_empty() {
114-
let ident_str = sig.ident.to_string();
118+
let ident_str = ident.to_string();
115119
if let Some(stripped) = strip_prefix(&ident_str, "tp_") {
116-
proc_macro2::Ident::new(stripped, sig.ident.span())
120+
proc_macro2::Ident::new(stripped, ident.span())
117121
} else {
118-
sig.ident.clone()
122+
ident.clone()
119123
}
120124
} else {
121125
match nesteds.into_iter().next().unwrap() {
@@ -128,14 +132,14 @@ impl Class {
128132
};
129133
Ok(ClassItem::Slot {
130134
slot_ident,
131-
item_ident: sig.ident.clone(),
135+
item_ident: ident.clone(),
132136
})
133137
}
134138

135139
fn extract_item_from_syn(
136140
&mut self,
137141
attrs: &mut Vec<Attribute>,
138-
sig: &Signature,
142+
ident: &Ident,
139143
) -> Result<(), Diagnostic> {
140144
let mut attr_idxs = Vec::new();
141145
for (i, meta) in attrs
@@ -149,10 +153,10 @@ impl Class {
149153
None => continue,
150154
};
151155
let item = match name.to_string().as_str() {
152-
"pymethod" => Self::extract_method(sig, meta)?,
153-
"pyclassmethod" => Self::extract_classmethod(sig, meta)?,
154-
"pyproperty" => Self::extract_property(sig, meta)?,
155-
"pyslot" => Self::extract_slot(sig, meta)?,
156+
"pymethod" => Self::extract_method(ident, meta)?,
157+
"pyclassmethod" => Self::extract_classmethod(ident, meta)?,
158+
"pyproperty" => Self::extract_property(ident, meta)?,
159+
"pyslot" => Self::extract_slot(ident, meta)?,
156160
_ => {
157161
continue;
158162
}
@@ -177,20 +181,15 @@ impl Class {
177181
}
178182
}
179183

180-
struct ItemSig<'a> {
181-
attrs: &'a mut Vec<Attribute>,
182-
sig: &'a Signature,
183-
}
184-
185-
fn extract_impl_items(mut items: Vec<ItemSig>) -> Result<TokenStream2, Diagnostic> {
184+
fn extract_impl_items(mut items: Vec<ItemIdent>) -> Result<TokenStream2, Diagnostic> {
186185
let mut diagnostics: Vec<Diagnostic> = Vec::new();
187186

188187
let mut class = Class::default();
189188

190189
for item in items.iter_mut() {
191190
push_diag_result!(
192191
diagnostics,
193-
class.extract_item_from_syn(&mut item.attrs, item.sig),
192+
class.extract_item_from_syn(&mut item.attrs, &item.ident),
194193
);
195194
}
196195

@@ -356,7 +355,10 @@ pub fn impl_pyimpl(attr: AttributeArgs, item: Item) -> Result<TokenStream2, Diag
356355
.iter_mut()
357356
.filter_map(|item| match item {
358357
syn::ImplItem::Method(syn::ImplItemMethod { attrs, sig, .. }) => {
359-
Some(ItemSig { attrs, sig })
358+
Some(ItemIdent {
359+
attrs,
360+
ident: &sig.ident,
361+
})
360362
}
361363
_ => None,
362364
})
@@ -386,7 +388,10 @@ pub fn impl_pyimpl(attr: AttributeArgs, item: Item) -> Result<TokenStream2, Diag
386388
.iter_mut()
387389
.filter_map(|item| match item {
388390
syn::TraitItem::Method(syn::TraitItemMethod { attrs, sig, .. }) => {
389-
Some(ItemSig { attrs, sig })
391+
Some(ItemIdent {
392+
attrs,
393+
ident: &sig.ident,
394+
})
390395
}
391396
_ => None,
392397
})

‎derive/src/pymodule.rs‎

Lines changed: 58 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,9 @@
11
use super::Diagnostic;
2-
use crate::util::{def_to_name, ItemMeta};
2+
use crate::util::{def_to_name, ItemIdent, ItemMeta};
33
use proc_macro2::{Span, TokenStream as TokenStream2};
44
use quote::{quote, quote_spanned};
55
use std::collections::HashSet;
6-
use syn::{
7-
parse_quote, spanned::Spanned, Attribute, AttributeArgs, Ident, Item, Meta, NestedMeta,
8-
Signature,
9-
};
6+
use syn::{parse_quote, spanned::Spanned, Attribute, AttributeArgs, Ident, Item, Meta, NestedMeta};
107

118
fn meta_to_vec(meta: Meta) -> Result<Vec<NestedMeta>, Meta> {
129
match meta {
@@ -24,6 +21,7 @@ struct Module {
2421
#[derive(PartialEq, Eq, Hash)]
2522
enum ModuleItem {
2623
Function { item_ident: Ident, py_name: String },
24+
Class { item_ident: Ident, py_name: String },
2725
}
2826

2927
impl Module {
@@ -38,7 +36,7 @@ impl Module {
3836
}
3937
}
4038

41-
fn extract_function(sig: &Signature, meta: Meta) -> Result<ModuleItem, Diagnostic> {
39+
fn extract_function(ident: &Ident, meta: Meta) -> Result<ModuleItem, Diagnostic> {
4240
let nesteds = meta_to_vec(meta).map_err(|meta| {
4341
err_span!(
4442
meta,
@@ -48,17 +46,34 @@ impl Module {
4846
})?;
4947

5048
let item_meta =
51-
ItemMeta::from_nested_meta("pyfunction", sig, &nesteds, ItemMeta::SIMPLE_NAMES)?;
49+
ItemMeta::from_nested_meta("pyfunction", &ident, &nesteds, ItemMeta::SIMPLE_NAMES)?;
5250
Ok(ModuleItem::Function {
53-
item_ident: sig.ident.clone(),
51+
item_ident: ident.clone(),
52+
py_name: item_meta.simple_name()?,
53+
})
54+
}
55+
56+
fn extract_class(ident: &Ident, meta: Meta) -> Result<ModuleItem, Diagnostic> {
57+
let nesteds = meta_to_vec(meta).map_err(|meta| {
58+
err_span!(
59+
meta,
60+
"#[pyclass = \"...\"] cannot be a name/value, you probably meant \
61+
#[pyclass(name = \"...\")]",
62+
)
63+
})?;
64+
65+
let item_meta =
66+
ItemMeta::from_nested_meta("pyclass", &ident, &nesteds, ItemMeta::SIMPLE_NAMES)?;
67+
Ok(ModuleItem::Class {
68+
item_ident: ident.clone(),
5469
py_name: item_meta.simple_name()?,
5570
})
5671
}
5772

5873
fn extract_item_from_syn(
5974
&mut self,
6075
attrs: &mut Vec<Attribute>,
61-
sig: &Signature,
76+
ident: &Ident,
6277
) -> Result<(), Diagnostic> {
6378
let mut attr_idxs = Vec::new();
6479
for (i, meta) in attrs
@@ -71,12 +86,17 @@ impl Module {
7186
Some(name) => name,
7287
None => continue,
7388
};
74-
if name == "pyfunction" {
75-
self.add_item(Self::extract_function(sig, meta)?, meta_span)?;
76-
} else {
77-
continue;
78-
}
79-
attr_idxs.push(i);
89+
let item = match name.to_string().as_str() {
90+
"pyfunction" => {
91+
attr_idxs.push(i);
92+
Self::extract_function(ident, meta)?
93+
}
94+
"pyclass" => Self::extract_class(ident, meta)?,
95+
_ => {
96+
continue;
97+
}
98+
};
99+
self.add_item(item, meta_span)?;
80100
}
81101
let mut i = 0;
82102
let mut attr_idxs = &*attr_idxs;
@@ -95,33 +115,37 @@ impl Module {
95115
}
96116
}
97117

98-
struct ItemSig<'a> {
99-
attrs: &'a mut Vec<Attribute>,
100-
sig: &'a Signature,
101-
}
102-
103-
fn extract_module_items(mut items: Vec<ItemSig>) -> Result<TokenStream2, Diagnostic> {
118+
fn extract_module_items(mut items: Vec<ItemIdent>) -> Result<TokenStream2, Diagnostic> {
104119
let mut diagnostics: Vec<Diagnostic> = Vec::new();
105120

106-
let mut class = Module::default();
121+
let mut module = Module::default();
107122

108123
for item in items.iter_mut() {
109124
push_diag_result!(
110125
diagnostics,
111-
class.extract_item_from_syn(&mut item.attrs, item.sig),
126+
module.extract_item_from_syn(&mut item.attrs, item.ident),
112127
);
113128
}
114129

115-
let functions = class.items.into_iter().map(|item| match item {
130+
let functions = module.items.into_iter().map(|item| match item {
116131
ModuleItem::Function {
117132
item_ident,
118133
py_name,
119134
} => {
120-
let new_func = quote_spanned!(item_ident.span()=> .new_function(#item_ident));
135+
let new_func = quote_spanned!(item_ident.span() => .new_function(#item_ident));
121136
quote! {
122137
vm.__module_set_attr(&module, #py_name, vm.ctx#new_func).unwrap();
123138
}
124139
}
140+
ModuleItem::Class {
141+
item_ident,
142+
py_name,
143+
} => {
144+
let new_class = quote_spanned!(item_ident.span() => #item_ident::make_class(&vm.ctx));
145+
quote! {
146+
vm.__module_set_attr(&module, #py_name, #new_class).unwrap();
147+
}
148+
}
125149
});
126150

127151
Diagnostic::from_vec(diagnostics)?;
@@ -140,12 +164,16 @@ pub fn impl_pymodule(attr: AttributeArgs, item: Item) -> Result<TokenStream2, Di
140164
let items = content
141165
.iter_mut()
142166
.filter_map(|item| match item {
143-
syn::Item::Fn(syn::ItemFn {
167+
Item::Fn(syn::ItemFn { attrs, sig, .. }) => Some(ItemIdent {
144168
attrs,
145-
vis: _vis,
146-
sig,
147-
..
148-
}) => Some(ItemSig { attrs, sig }),
169+
ident: &sig.ident,
170+
}),
171+
Item::Struct(syn::ItemStruct { attrs, ident, .. }) => {
172+
Some(ItemIdent { attrs, ident })
173+
}
174+
Item::Enum(syn::ItemEnum { attrs, ident, .. }) => {
175+
Some(ItemIdent { attrs, ident })
176+
}
149177
_ => None,
150178
})
151179
.collect();

0 commit comments

Comments
 (0)