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