1use std::collections::HashSet;
2
3use crate::utils::{has_attribute, has_attribute_with_namespace, Ctx, PyO3CratePath};
4use crate::{
5 attributes::{take_pyo3_options, CrateAttribute},
6 konst::{ConstAttributes, ConstSpec},
7 pyfunction::PyFunctionOptions,
8 pymethod::{self, is_proto_method, MethodAndMethodDef, MethodAndSlotDef},
9};
10use proc_macro2::TokenStream;
11use pymethod::GeneratedPyMethod;
12use quote::{format_ident, quote};
13use syn::ImplItemFn;
14use syn::{
15 parse::{Parse, ParseStream},
16 spanned::Spanned,
17 Result,
18};
19
20/// The mechanism used to collect `#[pymethods]` into the type object
21#[derive(Copy, Clone)]
22pub enum PyClassMethodsType {
23 Specialization,
24 Inventory,
25}
26
27enum PyImplPyO3Option {
28 Crate(CrateAttribute),
29}
30
31impl Parse for PyImplPyO3Option {
32 fn parse(input: ParseStream<'_>) -> Result<Self> {
33 let lookahead: Lookahead1<'_> = input.lookahead1();
34 if lookahead.peek(syn::Token![crate]) {
35 input.parse().map(op:PyImplPyO3Option::Crate)
36 } else {
37 Err(lookahead.error())
38 }
39 }
40}
41
42#[derive(Default)]
43pub struct PyImplOptions {
44 krate: Option<CrateAttribute>,
45}
46
47impl PyImplOptions {
48 pub fn from_attrs(attrs: &mut Vec<syn::Attribute>) -> Result<Self> {
49 let mut options: PyImplOptions = Default::default();
50
51 for option: PyImplPyO3Option in take_pyo3_options(attrs)? {
52 match option {
53 PyImplPyO3Option::Crate(path: KeywordAttribute>) => options.set_crate(path)?,
54 }
55 }
56
57 Ok(options)
58 }
59
60 fn set_crate(&mut self, path: CrateAttribute) -> Result<()> {
61 ensure_spanned!(
62 self.krate.is_none(),
63 path.span() => "`crate` may only be specified once"
64 );
65
66 self.krate = Some(path);
67 Ok(())
68 }
69}
70
71pub fn build_py_methods(
72 ast: &mut syn::ItemImpl,
73 methods_type: PyClassMethodsType,
74) -> syn::Result<TokenStream> {
75 if let Some((_, path: &Path, _)) = &ast.trait_ {
76 bail_spanned!(path.span() => "#[pymethods] cannot be used on trait impl blocks");
77 } else if ast.generics != Default::default() {
78 bail_spanned!(
79 ast.generics.span() =>
80 "#[pymethods] cannot be used with lifetime parameters or generics"
81 );
82 } else {
83 let options: PyImplOptions = PyImplOptions::from_attrs(&mut ast.attrs)?;
84 impl_methods(&ast.self_ty, &mut ast.items, methods_type, options)
85 }
86}
87
88fn check_pyfunction(pyo3_path: &PyO3CratePath, meth: &mut ImplItemFn) -> syn::Result<()> {
89 let mut error: Option = None;
90
91 meth.attrs.retain(|attr: &Attribute| {
92 let attrs: [Attribute; 1] = [attr.clone()];
93
94 if has_attribute(&attrs, ident:"pyfunction")
95 || has_attribute_with_namespace(&attrs, crate_path:Some(pyo3_path), &["pyfunction"])
96 || has_attribute_with_namespace(&attrs, crate_path:Some(pyo3_path), &["prelude", "pyfunction"]) {
97 error = Some(err_spanned!(meth.sig.span() => "functions inside #[pymethods] do not need to be annotated with #[pyfunction]"));
98 false
99 } else {
100 true
101 }
102 });
103
104 error.map_or(default:Ok(()), f:Err)
105}
106
107pub fn impl_methods(
108 ty: &syn::Type,
109 impls: &mut [syn::ImplItem],
110 methods_type: PyClassMethodsType,
111 options: PyImplOptions,
112) -> syn::Result<TokenStream> {
113 let mut trait_impls = Vec::new();
114 let mut proto_impls = Vec::new();
115 let mut methods = Vec::new();
116 let mut associated_methods = Vec::new();
117
118 let mut implemented_proto_fragments = HashSet::new();
119
120 for iimpl in impls {
121 match iimpl {
122 syn::ImplItem::Fn(meth) => {
123 let ctx = &Ctx::new(&options.krate, Some(&meth.sig));
124 let mut fun_options = PyFunctionOptions::from_attrs(&mut meth.attrs)?;
125 fun_options.krate = fun_options.krate.or_else(|| options.krate.clone());
126
127 check_pyfunction(&ctx.pyo3_path, meth)?;
128
129 match pymethod::gen_py_method(ty, &mut meth.sig, &mut meth.attrs, fun_options, ctx)?
130 {
131 GeneratedPyMethod::Method(MethodAndMethodDef {
132 associated_method,
133 method_def,
134 }) => {
135 let attrs = get_cfg_attributes(&meth.attrs);
136 associated_methods.push(quote!(#(#attrs)* #associated_method));
137 methods.push(quote!(#(#attrs)* #method_def));
138 }
139 GeneratedPyMethod::SlotTraitImpl(method_name, token_stream) => {
140 implemented_proto_fragments.insert(method_name);
141 let attrs = get_cfg_attributes(&meth.attrs);
142 trait_impls.push(quote!(#(#attrs)* #token_stream));
143 }
144 GeneratedPyMethod::Proto(MethodAndSlotDef {
145 associated_method,
146 slot_def,
147 }) => {
148 let attrs = get_cfg_attributes(&meth.attrs);
149 proto_impls.push(quote!(#(#attrs)* #slot_def));
150 associated_methods.push(quote!(#(#attrs)* #associated_method));
151 }
152 }
153 }
154 syn::ImplItem::Const(konst) => {
155 let ctx = &Ctx::new(&options.krate, None);
156 let attributes = ConstAttributes::from_attrs(&mut konst.attrs)?;
157 if attributes.is_class_attr {
158 let spec = ConstSpec {
159 rust_ident: konst.ident.clone(),
160 attributes,
161 };
162 let attrs = get_cfg_attributes(&konst.attrs);
163 let MethodAndMethodDef {
164 associated_method,
165 method_def,
166 } = gen_py_const(ty, &spec, ctx);
167 methods.push(quote!(#(#attrs)* #method_def));
168 associated_methods.push(quote!(#(#attrs)* #associated_method));
169 if is_proto_method(&spec.python_name().to_string()) {
170 // If this is a known protocol method e.g. __contains__, then allow this
171 // symbol even though it's not an uppercase constant.
172 konst
173 .attrs
174 .push(syn::parse_quote!(#[allow(non_upper_case_globals)]));
175 }
176 }
177 }
178 syn::ImplItem::Macro(m) => bail_spanned!(
179 m.span() =>
180 "macros cannot be used as items in `#[pymethods]` impl blocks\n\
181 = note: this was previously accepted and ignored"
182 ),
183 _ => {}
184 }
185 }
186 let ctx = &Ctx::new(&options.krate, None);
187
188 add_shared_proto_slots(ty, &mut proto_impls, implemented_proto_fragments, ctx);
189
190 let items = match methods_type {
191 PyClassMethodsType::Specialization => impl_py_methods(ty, methods, proto_impls, ctx),
192 PyClassMethodsType::Inventory => submit_methods_inventory(ty, methods, proto_impls, ctx),
193 };
194
195 Ok(quote! {
196 #(#trait_impls)*
197
198 #items
199
200 #[doc(hidden)]
201 #[allow(non_snake_case)]
202 impl #ty {
203 #(#associated_methods)*
204 }
205 })
206}
207
208pub fn gen_py_const(cls: &syn::Type, spec: &ConstSpec, ctx: &Ctx) -> MethodAndMethodDef {
209 let member = &spec.rust_ident;
210 let wrapper_ident = format_ident!("__pymethod_{}__", member);
211 let python_name = spec.null_terminated_python_name(ctx);
212 let Ctx { pyo3_path, .. } = ctx;
213
214 let associated_method = quote! {
215 fn #wrapper_ident(py: #pyo3_path::Python<'_>) -> #pyo3_path::PyResult<#pyo3_path::PyObject> {
216 #pyo3_path::IntoPyObjectExt::into_py_any(#cls::#member, py)
217 }
218 };
219
220 let method_def = quote! {
221 #pyo3_path::impl_::pyclass::MaybeRuntimePyMethodDef::Static(
222 #pyo3_path::impl_::pymethods::PyMethodDefType::ClassAttribute({
223 #pyo3_path::impl_::pymethods::PyClassAttributeDef::new(
224 #python_name,
225 #cls::#wrapper_ident
226 )
227 })
228 )
229 };
230
231 MethodAndMethodDef {
232 associated_method,
233 method_def,
234 }
235}
236
237fn impl_py_methods(
238 ty: &syn::Type,
239 methods: Vec<TokenStream>,
240 proto_impls: Vec<TokenStream>,
241 ctx: &Ctx,
242) -> TokenStream {
243 let Ctx { pyo3_path: &PyO3CratePath, .. } = ctx;
244 quote! {
245 #[allow(unknown_lints, non_local_definitions)]
246 impl #pyo3_path::impl_::pyclass::PyMethods<#ty>
247 for #pyo3_path::impl_::pyclass::PyClassImplCollector<#ty>
248 {
249 fn py_methods(self) -> &'static #pyo3_path::impl_::pyclass::PyClassItems {
250 static ITEMS: #pyo3_path::impl_::pyclass::PyClassItems = #pyo3_path::impl_::pyclass::PyClassItems {
251 methods: &[#(#methods),*],
252 slots: &[#(#proto_impls),*]
253 };
254 &ITEMS
255 }
256 }
257 }
258}
259
260fn add_shared_proto_slots(
261 ty: &syn::Type,
262 proto_impls: &mut Vec<TokenStream>,
263 mut implemented_proto_fragments: HashSet<String>,
264 ctx: &Ctx,
265) {
266 let Ctx { pyo3_path, .. } = ctx;
267 macro_rules! try_add_shared_slot {
268 ($slot:ident, $($fragments:literal),*) => {{
269 let mut implemented = false;
270 $(implemented |= implemented_proto_fragments.remove($fragments));*;
271 if implemented {
272 proto_impls.push(quote! { #pyo3_path::impl_::pyclass::$slot!(#ty) })
273 }
274 }};
275 }
276
277 try_add_shared_slot!(
278 generate_pyclass_getattro_slot,
279 "__getattribute__",
280 "__getattr__"
281 );
282 try_add_shared_slot!(generate_pyclass_setattr_slot, "__setattr__", "__delattr__");
283 try_add_shared_slot!(generate_pyclass_setdescr_slot, "__set__", "__delete__");
284 try_add_shared_slot!(generate_pyclass_setitem_slot, "__setitem__", "__delitem__");
285 try_add_shared_slot!(generate_pyclass_add_slot, "__add__", "__radd__");
286 try_add_shared_slot!(generate_pyclass_sub_slot, "__sub__", "__rsub__");
287 try_add_shared_slot!(generate_pyclass_mul_slot, "__mul__", "__rmul__");
288 try_add_shared_slot!(generate_pyclass_mod_slot, "__mod__", "__rmod__");
289 try_add_shared_slot!(generate_pyclass_divmod_slot, "__divmod__", "__rdivmod__");
290 try_add_shared_slot!(generate_pyclass_lshift_slot, "__lshift__", "__rlshift__");
291 try_add_shared_slot!(generate_pyclass_rshift_slot, "__rshift__", "__rrshift__");
292 try_add_shared_slot!(generate_pyclass_and_slot, "__and__", "__rand__");
293 try_add_shared_slot!(generate_pyclass_or_slot, "__or__", "__ror__");
294 try_add_shared_slot!(generate_pyclass_xor_slot, "__xor__", "__rxor__");
295 try_add_shared_slot!(generate_pyclass_matmul_slot, "__matmul__", "__rmatmul__");
296 try_add_shared_slot!(generate_pyclass_truediv_slot, "__truediv__", "__rtruediv__");
297 try_add_shared_slot!(
298 generate_pyclass_floordiv_slot,
299 "__floordiv__",
300 "__rfloordiv__"
301 );
302 try_add_shared_slot!(generate_pyclass_pow_slot, "__pow__", "__rpow__");
303 try_add_shared_slot!(
304 generate_pyclass_richcompare_slot,
305 "__lt__",
306 "__le__",
307 "__eq__",
308 "__ne__",
309 "__gt__",
310 "__ge__"
311 );
312
313 // if this assertion trips, a slot fragment has been implemented which has not been added in the
314 // list above
315 assert!(implemented_proto_fragments.is_empty());
316}
317
318fn submit_methods_inventory(
319 ty: &syn::Type,
320 methods: Vec<TokenStream>,
321 proto_impls: Vec<TokenStream>,
322 ctx: &Ctx,
323) -> TokenStream {
324 let Ctx { pyo3_path: &PyO3CratePath, .. } = ctx;
325 quote! {
326 #pyo3_path::inventory::submit! {
327 type Inventory = <#ty as #pyo3_path::impl_::pyclass::PyClassImpl>::Inventory;
328 Inventory::new(#pyo3_path::impl_::pyclass::PyClassItems { methods: &[#(#methods),*], slots: &[#(#proto_impls),*] })
329 }
330 }
331}
332
333pub(crate) fn get_cfg_attributes(attrs: &[syn::Attribute]) -> Vec<&syn::Attribute> {
334 attrsimpl Iterator
335 .iter()
336 .filter(|attr: &&Attribute| attr.path().is_ident("cfg"))
337 .collect()
338}
339