1use crate::utils::{Ctx, TypeExt as _};
2use crate::{
3 attributes::FromPyWithAttribute,
4 method::{FnArg, FnSpec, RegularArg},
5 pyfunction::FunctionSignature,
6 quotes::some_wrap,
7};
8use proc_macro2::{Span, TokenStream};
9use quote::{format_ident, quote, quote_spanned};
10use syn::spanned::Spanned;
11
12pub struct Holders {
13 holders: Vec<syn::Ident>,
14}
15
16impl Holders {
17 pub fn new() -> Self {
18 Holders {
19 holders: Vec::new(),
20 }
21 }
22
23 pub fn push_holder(&mut self, span: Span) -> syn::Ident {
24 let holder = syn::Ident::new(&format!("holder_{}", self.holders.len()), span);
25 self.holders.push(holder.clone());
26 holder
27 }
28
29 pub fn init_holders(&self, ctx: &Ctx) -> TokenStream {
30 let Ctx { pyo3_path, .. } = ctx;
31 let holders = &self.holders;
32 quote! {
33 #[allow(clippy::let_unit_value)]
34 #(let mut #holders = #pyo3_path::impl_::extract_argument::FunctionArgumentHolder::INIT;)*
35 }
36 }
37}
38
39pub fn is_forwarded_args(signature: &FunctionSignature<'_>) -> bool {
41 matches!(
42 signature.arguments.as_slice(),
43 [FnArg::VarArgs(..), FnArg::KwArgs(..),]
44 )
45}
46
47pub fn impl_arg_params(
48 spec: &FnSpec<'_>,
49 self_: Option<&syn::Type>,
50 fastcall: bool,
51 holders: &mut Holders,
52 ctx: &Ctx,
53) -> (TokenStream, Vec<TokenStream>) {
54 let args_array = syn::Ident::new("output", Span::call_site());
55 let Ctx { pyo3_path, .. } = ctx;
56
57 let from_py_with = spec
58 .signature
59 .arguments
60 .iter()
61 .enumerate()
62 .filter_map(|(i, arg)| {
63 let from_py_with = &arg.from_py_with()?.value;
64 let from_py_with_holder = format_ident!("from_py_with_{}", i);
65 Some(quote_spanned! { from_py_with.span() =>
66 let #from_py_with_holder = #from_py_with;
67 })
68 })
69 .collect::<TokenStream>();
70
71 if !fastcall && is_forwarded_args(&spec.signature) {
72 let arg_convert = spec
75 .signature
76 .arguments
77 .iter()
78 .enumerate()
79 .map(|(i, arg)| impl_arg_param(arg, i, &mut 0, holders, ctx))
80 .collect();
81 return (
82 quote! {
83 let _args = unsafe { #pyo3_path::impl_::pymethods::BoundRef::ref_from_ptr(py, &_args) };
84 let _kwargs = #pyo3_path::impl_::pymethods::BoundRef::ref_from_ptr_or_opt(py, &_kwargs);
85 #from_py_with
86 },
87 arg_convert,
88 );
89 };
90
91 let positional_parameter_names = &spec.signature.python_signature.positional_parameters;
92 let positional_only_parameters = &spec.signature.python_signature.positional_only_parameters;
93 let required_positional_parameters = &spec
94 .signature
95 .python_signature
96 .required_positional_parameters;
97 let keyword_only_parameters = spec
98 .signature
99 .python_signature
100 .keyword_only_parameters
101 .iter()
102 .map(|(name, required)| {
103 quote! {
104 #pyo3_path::impl_::extract_argument::KeywordOnlyParameterDescription {
105 name: #name,
106 required: #required,
107 }
108 }
109 });
110
111 let num_params = positional_parameter_names.len() + keyword_only_parameters.len();
112
113 let mut option_pos = 0usize;
114 let param_conversion = spec
115 .signature
116 .arguments
117 .iter()
118 .enumerate()
119 .map(|(i, arg)| impl_arg_param(arg, i, &mut option_pos, holders, ctx))
120 .collect();
121
122 let args_handler = if spec.signature.python_signature.varargs.is_some() {
123 quote! { #pyo3_path::impl_::extract_argument::TupleVarargs }
124 } else {
125 quote! { #pyo3_path::impl_::extract_argument::NoVarargs }
126 };
127 let kwargs_handler = if spec.signature.python_signature.kwargs.is_some() {
128 quote! { #pyo3_path::impl_::extract_argument::DictVarkeywords }
129 } else {
130 quote! { #pyo3_path::impl_::extract_argument::NoVarkeywords }
131 };
132
133 let cls_name = if let Some(cls) = self_ {
134 quote! { ::std::option::Option::Some(<#cls as #pyo3_path::type_object::PyTypeInfo>::NAME) }
135 } else {
136 quote! { ::std::option::Option::None }
137 };
138 let python_name = &spec.python_name;
139
140 let extract_expression = if fastcall {
141 quote! {
142 DESCRIPTION.extract_arguments_fastcall::<#args_handler, #kwargs_handler>(
143 py,
144 _args,
145 _nargs,
146 _kwnames,
147 &mut #args_array
148 )?
149 }
150 } else {
151 quote! {
152 DESCRIPTION.extract_arguments_tuple_dict::<#args_handler, #kwargs_handler>(
153 py,
154 _args,
155 _kwargs,
156 &mut #args_array
157 )?
158 }
159 };
160
161 (
163 quote! {
164 const DESCRIPTION: #pyo3_path::impl_::extract_argument::FunctionDescription = #pyo3_path::impl_::extract_argument::FunctionDescription {
165 cls_name: #cls_name,
166 func_name: stringify!(#python_name),
167 positional_parameter_names: &[#(#positional_parameter_names),*],
168 positional_only_parameters: #positional_only_parameters,
169 required_positional_parameters: #required_positional_parameters,
170 keyword_only_parameters: &[#(#keyword_only_parameters),*],
171 };
172 let mut #args_array = [::std::option::Option::None; #num_params];
173 let (_args, _kwargs) = #extract_expression;
174 #from_py_with
175 },
176 param_conversion,
177 )
178}
179
180fn impl_arg_param(
181 arg: &FnArg<'_>,
182 pos: usize,
183 option_pos: &mut usize,
184 holders: &mut Holders,
185 ctx: &Ctx,
186) -> TokenStream {
187 let Ctx { pyo3_path, .. } = ctx;
188 let args_array = syn::Ident::new("output", Span::call_site());
189
190 match arg {
191 FnArg::Regular(arg) => {
192 let from_py_with = format_ident!("from_py_with_{}", pos);
193 let arg_value = quote!(#args_array[#option_pos].as_deref());
194 *option_pos += 1;
195 impl_regular_arg_param(arg, from_py_with, arg_value, holders, ctx)
196 }
197 FnArg::VarArgs(arg) => {
198 let holder = holders.push_holder(arg.name.span());
199 let name_str = arg.name.to_string();
200 quote_spanned! { arg.name.span() =>
201 #pyo3_path::impl_::extract_argument::extract_argument::<_, false>(
202 &_args,
203 &mut #holder,
204 #name_str
205 )?
206 }
207 }
208 FnArg::KwArgs(arg) => {
209 let holder = holders.push_holder(arg.name.span());
210 let name_str = arg.name.to_string();
211 quote_spanned! { arg.name.span() =>
212 #pyo3_path::impl_::extract_argument::extract_optional_argument::<_, false>(
213 _kwargs.as_deref(),
214 &mut #holder,
215 #name_str,
216 || ::std::option::Option::None
217 )?
218 }
219 }
220 FnArg::Py(..) => quote! { py },
221 FnArg::CancelHandle(..) => quote! { __cancel_handle },
222 }
223}
224
225pub(crate) fn impl_regular_arg_param(
228 arg: &RegularArg<'_>,
229 from_py_with: syn::Ident,
230 arg_value: TokenStream, holders: &mut Holders,
232 ctx: &Ctx,
233) -> TokenStream {
234 let Ctx { pyo3_path, .. } = ctx;
235 let pyo3_path = pyo3_path.to_tokens_spanned(arg.ty.span());
236
237 let use_probe = quote! {
240 #[allow(unused_imports)]
241 use #pyo3_path::impl_::pyclass::Probe as _;
242 };
243 macro_rules! quote_arg_span {
244 ($($tokens:tt)*) => { quote_spanned!(arg.ty.span() => { #use_probe $($tokens)* }) }
245 }
246
247 let name_str = arg.name.to_string();
248 let mut default = arg.default_value.as_ref().map(|expr| quote!(#expr));
249
250 if arg.option_wrapped_type.is_some() {
253 default = default.map(|tokens| some_wrap(tokens, ctx));
254 }
255
256 let arg_ty = arg.ty.clone().elide_lifetimes();
257 if let Some(FromPyWithAttribute { kw, .. }) = arg.from_py_with {
258 let extractor = quote_spanned! { kw.span =>
259 { let from_py_with: fn(_) -> _ = #from_py_with; from_py_with }
260 };
261 if let Some(default) = default {
262 quote_arg_span! {
263 #pyo3_path::impl_::extract_argument::from_py_with_with_default(
264 #arg_value,
265 #name_str,
266 #extractor,
267 #[allow(clippy::redundant_closure)]
268 {
269 || #default
270 }
271 )?
272 }
273 } else {
274 let unwrap = quote! {unsafe { #pyo3_path::impl_::extract_argument::unwrap_required_argument(#arg_value) }};
275 quote_arg_span! {
276 #pyo3_path::impl_::extract_argument::from_py_with(
277 #unwrap,
278 #name_str,
279 #extractor,
280 )?
281 }
282 }
283 } else if let Some(default) = default {
284 let holder = holders.push_holder(arg.name.span());
285 if let Some(arg_ty) = arg.option_wrapped_type {
286 let arg_ty = arg_ty.clone().elide_lifetimes();
287 quote_arg_span! {
288 #pyo3_path::impl_::extract_argument::extract_optional_argument::<
289 _,
290 { #pyo3_path::impl_::pyclass::IsOption::<#arg_ty>::VALUE }
291 >(
292 #arg_value,
293 &mut #holder,
294 #name_str,
295 #[allow(clippy::redundant_closure)]
296 {
297 || #default
298 }
299 )?
300 }
301 } else {
302 quote_arg_span! {
303 #pyo3_path::impl_::extract_argument::extract_argument_with_default::<
304 _,
305 { #pyo3_path::impl_::pyclass::IsOption::<#arg_ty>::VALUE }
306 >(
307 #arg_value,
308 &mut #holder,
309 #name_str,
310 #[allow(clippy::redundant_closure)]
311 {
312 || #default
313 }
314 )?
315 }
316 }
317 } else {
318 let holder = holders.push_holder(arg.name.span());
319 let unwrap = quote! {unsafe { #pyo3_path::impl_::extract_argument::unwrap_required_argument(#arg_value) }};
320 quote_arg_span! {
321 #pyo3_path::impl_::extract_argument::extract_argument::<
322 _,
323 { #pyo3_path::impl_::pyclass::IsOption::<#arg_ty>::VALUE }
324 >(
325 #unwrap,
326 &mut #holder,
327 #name_str
328 )?
329 }
330 }
331}