1use proc_macro2::{Span, TokenStream};
2use quote::ToTokens;
3use syn::{
4 ext::IdentExt,
5 parse::{Parse, ParseStream},
6 punctuated::Punctuated,
7 spanned::Spanned,
8 Token,
9};
10
11use crate::{
12 attributes::{kw, KeywordAttribute},
13 method::{FnArg, RegularArg},
14};
15
16#[derive(Clone)]
17pub struct Signature {
18 paren_token: syn::token::Paren,
19 pub items: Punctuated<SignatureItem, Token![,]>,
20}
21
22impl Parse for Signature {
23 fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
24 let content: ParseBuffer<'_>;
25 let paren_token: Paren = syn::parenthesized!(content in input);
26
27 let items: Punctuated = content.parse_terminated(parser:SignatureItem::parse, separator:Token![,])?;
28
29 Ok(Signature { paren_token, items })
30 }
31}
32
33impl ToTokens for Signature {
34 fn to_tokens(&self, tokens: &mut TokenStream) {
35 self.paren_token
36 .surround(tokens, |tokens: &mut TokenStream| self.items.to_tokens(tokens))
37 }
38}
39
40#[derive(Clone, Debug, PartialEq, Eq)]
41pub struct SignatureItemArgument {
42 pub ident: syn::Ident,
43 pub eq_and_default: Option<(Token![=], syn::Expr)>,
44}
45
46#[derive(Clone, Debug, PartialEq, Eq)]
47pub struct SignatureItemPosargsSep {
48 pub slash: Token![/],
49}
50
51#[derive(Clone, Debug, PartialEq, Eq)]
52pub struct SignatureItemVarargsSep {
53 pub asterisk: Token![*],
54}
55
56#[derive(Clone, Debug, PartialEq, Eq)]
57pub struct SignatureItemVarargs {
58 pub sep: SignatureItemVarargsSep,
59 pub ident: syn::Ident,
60}
61
62#[derive(Clone, Debug, PartialEq, Eq)]
63pub struct SignatureItemKwargs {
64 pub asterisks: (Token![*], Token![*]),
65 pub ident: syn::Ident,
66}
67
68#[derive(Clone, Debug, PartialEq, Eq)]
69pub enum SignatureItem {
70 Argument(Box<SignatureItemArgument>),
71 PosargsSep(SignatureItemPosargsSep),
72 VarargsSep(SignatureItemVarargsSep),
73 Varargs(SignatureItemVarargs),
74 Kwargs(SignatureItemKwargs),
75}
76
77impl Parse for SignatureItem {
78 fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
79 let lookahead: Lookahead1<'_> = input.lookahead1();
80 if lookahead.peek(Token![*]) {
81 if input.peek2(Token![*]) {
82 input.parse().map(op:SignatureItem::Kwargs)
83 } else {
84 let sep: SignatureItemVarargsSep = input.parse()?;
85 if input.is_empty() || input.peek(Token![,]) {
86 Ok(SignatureItem::VarargsSep(sep))
87 } else {
88 Ok(SignatureItem::Varargs(SignatureItemVarargs {
89 sep,
90 ident: input.parse()?,
91 }))
92 }
93 }
94 } else if lookahead.peek(Token![/]) {
95 input.parse().map(op:SignatureItem::PosargsSep)
96 } else {
97 input.parse().map(op:SignatureItem::Argument)
98 }
99 }
100}
101
102impl ToTokens for SignatureItem {
103 fn to_tokens(&self, tokens: &mut TokenStream) {
104 match self {
105 SignatureItem::Argument(arg: &Box) => arg.to_tokens(tokens),
106 SignatureItem::Varargs(varargs: &SignatureItemVarargs) => varargs.to_tokens(tokens),
107 SignatureItem::VarargsSep(sep: &SignatureItemVarargsSep) => sep.to_tokens(tokens),
108 SignatureItem::Kwargs(kwargs: &SignatureItemKwargs) => kwargs.to_tokens(tokens),
109 SignatureItem::PosargsSep(sep: &SignatureItemPosargsSep) => sep.to_tokens(tokens),
110 }
111 }
112}
113
114impl Parse for SignatureItemArgument {
115 fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
116 Ok(Self {
117 ident: input.parse()?,
118 eq_and_default: if input.peek(Token![=]) {
119 Some((input.parse()?, input.parse()?))
120 } else {
121 None
122 },
123 })
124 }
125}
126
127impl ToTokens for SignatureItemArgument {
128 fn to_tokens(&self, tokens: &mut TokenStream) {
129 self.ident.to_tokens(tokens);
130 if let Some((eq: &Eq, default: &Expr)) = &self.eq_and_default {
131 eq.to_tokens(tokens);
132 default.to_tokens(tokens);
133 }
134 }
135}
136
137impl Parse for SignatureItemVarargsSep {
138 fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
139 Ok(Self {
140 asterisk: input.parse()?,
141 })
142 }
143}
144
145impl ToTokens for SignatureItemVarargsSep {
146 fn to_tokens(&self, tokens: &mut TokenStream) {
147 self.asterisk.to_tokens(tokens);
148 }
149}
150
151impl Parse for SignatureItemVarargs {
152 fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
153 Ok(Self {
154 sep: input.parse()?,
155 ident: input.parse()?,
156 })
157 }
158}
159
160impl ToTokens for SignatureItemVarargs {
161 fn to_tokens(&self, tokens: &mut TokenStream) {
162 self.sep.to_tokens(tokens);
163 self.ident.to_tokens(tokens);
164 }
165}
166
167impl Parse for SignatureItemKwargs {
168 fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
169 Ok(Self {
170 asterisks: (input.parse()?, input.parse()?),
171 ident: input.parse()?,
172 })
173 }
174}
175
176impl ToTokens for SignatureItemKwargs {
177 fn to_tokens(&self, tokens: &mut TokenStream) {
178 self.asterisks.0.to_tokens(tokens);
179 self.asterisks.1.to_tokens(tokens);
180 self.ident.to_tokens(tokens);
181 }
182}
183
184impl Parse for SignatureItemPosargsSep {
185 fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
186 Ok(Self {
187 slash: input.parse()?,
188 })
189 }
190}
191
192impl ToTokens for SignatureItemPosargsSep {
193 fn to_tokens(&self, tokens: &mut TokenStream) {
194 self.slash.to_tokens(tokens);
195 }
196}
197
198pub type SignatureAttribute = KeywordAttribute<kw::signature, Signature>;
199pub type ConstructorAttribute = KeywordAttribute<kw::constructor, Signature>;
200
201impl ConstructorAttribute {
202 pub fn into_signature(self) -> SignatureAttribute {
203 SignatureAttribute {
204 kw: kw::signature(self.kw.span),
205 value: self.value,
206 }
207 }
208}
209
210#[derive(Default)]
211pub struct PythonSignature {
212 pub positional_parameters: Vec<String>,
213 pub positional_only_parameters: usize,
214 pub required_positional_parameters: usize,
215 pub varargs: Option<String>,
216 // Tuples of keyword name and whether it is required
217 pub keyword_only_parameters: Vec<(String, bool)>,
218 pub kwargs: Option<String>,
219}
220
221impl PythonSignature {
222 pub fn has_no_args(&self) -> bool {
223 self.positional_parameters.is_empty()
224 && self.keyword_only_parameters.is_empty()
225 && self.varargs.is_none()
226 && self.kwargs.is_none()
227 }
228}
229
230pub struct FunctionSignature<'a> {
231 pub arguments: Vec<FnArg<'a>>,
232 pub python_signature: PythonSignature,
233 pub attribute: Option<SignatureAttribute>,
234}
235
236pub enum ParseState {
237 /// Accepting positional parameters, which might be positional only
238 Positional,
239 /// Accepting positional parameters after '/'
240 PositionalAfterPosargs,
241 /// Accepting keyword-only parameters after '*' or '*args'
242 Keywords,
243 /// After `**kwargs` nothing is allowed
244 Done,
245}
246
247impl ParseState {
248 fn add_argument(
249 &mut self,
250 signature: &mut PythonSignature,
251 name: String,
252 required: bool,
253 span: Span,
254 ) -> syn::Result<()> {
255 match self {
256 ParseState::Positional | ParseState::PositionalAfterPosargs => {
257 signature.positional_parameters.push(name);
258 if required {
259 signature.required_positional_parameters += 1;
260 ensure_spanned!(
261 signature.required_positional_parameters == signature.positional_parameters.len(),
262 span => "cannot have required positional parameter after an optional parameter"
263 );
264 }
265 Ok(())
266 }
267 ParseState::Keywords => {
268 signature.keyword_only_parameters.push((name, required));
269 Ok(())
270 }
271 ParseState::Done => {
272 bail_spanned!(span => format!("no more arguments are allowed after `**{}`", signature.kwargs.as_deref().unwrap_or("")))
273 }
274 }
275 }
276
277 fn add_varargs(
278 &mut self,
279 signature: &mut PythonSignature,
280 varargs: &SignatureItemVarargs,
281 ) -> syn::Result<()> {
282 match self {
283 ParseState::Positional | ParseState::PositionalAfterPosargs => {
284 signature.varargs = Some(varargs.ident.to_string());
285 *self = ParseState::Keywords;
286 Ok(())
287 }
288 ParseState::Keywords => {
289 bail_spanned!(varargs.span() => format!("`*{}` not allowed after `*{}`", varargs.ident, signature.varargs.as_deref().unwrap_or("")))
290 }
291 ParseState::Done => {
292 bail_spanned!(varargs.span() => format!("`*{}` not allowed after `**{}`", varargs.ident, signature.kwargs.as_deref().unwrap_or("")))
293 }
294 }
295 }
296
297 fn add_kwargs(
298 &mut self,
299 signature: &mut PythonSignature,
300 kwargs: &SignatureItemKwargs,
301 ) -> syn::Result<()> {
302 match self {
303 ParseState::Positional | ParseState::PositionalAfterPosargs | ParseState::Keywords => {
304 signature.kwargs = Some(kwargs.ident.to_string());
305 *self = ParseState::Done;
306 Ok(())
307 }
308 ParseState::Done => {
309 bail_spanned!(kwargs.span() => format!("`**{}` not allowed after `**{}`", kwargs.ident, signature.kwargs.as_deref().unwrap_or("")))
310 }
311 }
312 }
313
314 fn finish_pos_only_args(
315 &mut self,
316 signature: &mut PythonSignature,
317 span: Span,
318 ) -> syn::Result<()> {
319 match self {
320 ParseState::Positional => {
321 signature.positional_only_parameters = signature.positional_parameters.len();
322 *self = ParseState::PositionalAfterPosargs;
323 Ok(())
324 }
325 ParseState::PositionalAfterPosargs => {
326 bail_spanned!(span => "`/` not allowed after `/`")
327 }
328 ParseState::Keywords => {
329 bail_spanned!(span => format!("`/` not allowed after `*{}`", signature.varargs.as_deref().unwrap_or("")))
330 }
331 ParseState::Done => {
332 bail_spanned!(span => format!("`/` not allowed after `**{}`", signature.kwargs.as_deref().unwrap_or("")))
333 }
334 }
335 }
336
337 fn finish_pos_args(&mut self, signature: &PythonSignature, span: Span) -> syn::Result<()> {
338 match self {
339 ParseState::Positional | ParseState::PositionalAfterPosargs => {
340 *self = ParseState::Keywords;
341 Ok(())
342 }
343 ParseState::Keywords => {
344 bail_spanned!(span => format!("`*` not allowed after `*{}`", signature.varargs.as_deref().unwrap_or("")))
345 }
346 ParseState::Done => {
347 bail_spanned!(span => format!("`*` not allowed after `**{}`", signature.kwargs.as_deref().unwrap_or("")))
348 }
349 }
350 }
351}
352
353impl<'a> FunctionSignature<'a> {
354 pub fn from_arguments_and_attribute(
355 mut arguments: Vec<FnArg<'a>>,
356 attribute: SignatureAttribute,
357 ) -> syn::Result<Self> {
358 let mut parse_state = ParseState::Positional;
359 let mut python_signature = PythonSignature::default();
360
361 let mut args_iter = arguments.iter_mut();
362
363 let mut next_non_py_argument_checked = |name: &syn::Ident| {
364 for fn_arg in args_iter.by_ref() {
365 match fn_arg {
366 crate::method::FnArg::Py(..) => {
367 // If the user incorrectly tried to include py: Python in the
368 // signature, give a useful error as a hint.
369 ensure_spanned!(
370 name != fn_arg.name(),
371 name.span() => "arguments of type `Python` must not be part of the signature"
372 );
373 // Otherwise try next argument.
374 continue;
375 }
376 crate::method::FnArg::CancelHandle(..) => {
377 // If the user incorrectly tried to include cancel: CoroutineCancel in the
378 // signature, give a useful error as a hint.
379 ensure_spanned!(
380 name != fn_arg.name(),
381 name.span() => "`cancel_handle` argument must not be part of the signature"
382 );
383 // Otherwise try next argument.
384 continue;
385 }
386 _ => {
387 ensure_spanned!(
388 name == fn_arg.name(),
389 name.span() => format!(
390 "expected argument from function definition `{}` but got argument `{}`",
391 fn_arg.name().unraw(),
392 name.unraw(),
393 )
394 );
395 return Ok(fn_arg);
396 }
397 }
398 }
399 bail_spanned!(
400 name.span() => "signature entry does not have a corresponding function argument"
401 )
402 };
403
404 for item in &attribute.value.items {
405 match item {
406 SignatureItem::Argument(arg) => {
407 let fn_arg = next_non_py_argument_checked(&arg.ident)?;
408 parse_state.add_argument(
409 &mut python_signature,
410 arg.ident.unraw().to_string(),
411 arg.eq_and_default.is_none(),
412 arg.span(),
413 )?;
414 if let Some((_, default)) = &arg.eq_and_default {
415 if let FnArg::Regular(arg) = fn_arg {
416 arg.default_value = Some(default.clone());
417 } else {
418 unreachable!(
419 "`Python` and `CancelHandle` are already handled above and `*args`/`**kwargs` are \
420 parsed and transformed below. Because the have to come last and are only allowed \
421 once, this has to be a regular argument."
422 );
423 }
424 }
425 }
426 SignatureItem::VarargsSep(sep) => {
427 parse_state.finish_pos_args(&python_signature, sep.span())?
428 }
429 SignatureItem::Varargs(varargs) => {
430 let fn_arg = next_non_py_argument_checked(&varargs.ident)?;
431 fn_arg.to_varargs_mut()?;
432 parse_state.add_varargs(&mut python_signature, varargs)?;
433 }
434 SignatureItem::Kwargs(kwargs) => {
435 let fn_arg = next_non_py_argument_checked(&kwargs.ident)?;
436 fn_arg.to_kwargs_mut()?;
437 parse_state.add_kwargs(&mut python_signature, kwargs)?;
438 }
439 SignatureItem::PosargsSep(sep) => {
440 parse_state.finish_pos_only_args(&mut python_signature, sep.span())?
441 }
442 };
443 }
444
445 // Ensure no non-py arguments remain
446 if let Some(arg) =
447 args_iter.find(|arg| !matches!(arg, FnArg::Py(..) | FnArg::CancelHandle(..)))
448 {
449 bail_spanned!(
450 attribute.kw.span() => format!("missing signature entry for argument `{}`", arg.name())
451 );
452 }
453
454 Ok(FunctionSignature {
455 arguments,
456 python_signature,
457 attribute: Some(attribute),
458 })
459 }
460
461 /// Without `#[pyo3(signature)]` or `#[args]` - just take the Rust function arguments as positional.
462 pub fn from_arguments(arguments: Vec<FnArg<'a>>) -> Self {
463 let mut python_signature = PythonSignature::default();
464 for arg in &arguments {
465 // Python<'_> arguments don't show in Python signature
466 if matches!(arg, FnArg::Py(..) | FnArg::CancelHandle(..)) {
467 continue;
468 }
469
470 if let FnArg::Regular(RegularArg { .. }) = arg {
471 // This argument is required, all previous arguments must also have been required
472 assert_eq!(
473 python_signature.required_positional_parameters,
474 python_signature.positional_parameters.len(),
475 );
476
477 python_signature.required_positional_parameters =
478 python_signature.positional_parameters.len() + 1;
479 }
480
481 python_signature
482 .positional_parameters
483 .push(arg.name().unraw().to_string());
484 }
485
486 Self {
487 arguments,
488 python_signature,
489 attribute: None,
490 }
491 }
492
493 fn default_value_for_parameter(&self, parameter: &str) -> String {
494 let mut default = "...".to_string();
495 if let Some(fn_arg) = self.arguments.iter().find(|arg| arg.name() == parameter) {
496 if let FnArg::Regular(RegularArg {
497 default_value: Some(arg_default),
498 ..
499 }) = fn_arg
500 {
501 match arg_default {
502 // literal values
503 syn::Expr::Lit(syn::ExprLit { lit, .. }) => match lit {
504 syn::Lit::Str(s) => default = s.token().to_string(),
505 syn::Lit::Char(c) => default = c.token().to_string(),
506 syn::Lit::Int(i) => default = i.base10_digits().to_string(),
507 syn::Lit::Float(f) => default = f.base10_digits().to_string(),
508 syn::Lit::Bool(b) => {
509 default = if b.value() {
510 "True".to_string()
511 } else {
512 "False".to_string()
513 }
514 }
515 _ => {}
516 },
517 // None
518 syn::Expr::Path(syn::ExprPath {
519 qself: None, path, ..
520 }) if path.is_ident("None") => {
521 default = "None".to_string();
522 }
523 // others, unsupported yet so defaults to `...`
524 _ => {}
525 }
526 } else if let FnArg::Regular(RegularArg {
527 option_wrapped_type: Some(..),
528 ..
529 }) = fn_arg
530 {
531 // functions without a `#[pyo3(signature = (...))]` option
532 // will treat trailing `Option<T>` arguments as having a default of `None`
533 default = "None".to_string();
534 }
535 }
536 default
537 }
538
539 pub fn text_signature(&self, self_argument: Option<&str>) -> String {
540 let mut output = String::new();
541 output.push('(');
542
543 if let Some(arg) = self_argument {
544 output.push('$');
545 output.push_str(arg);
546 }
547
548 let mut maybe_push_comma = {
549 let mut first = self_argument.is_none();
550 move |output: &mut String| {
551 if !first {
552 output.push_str(", ");
553 } else {
554 first = false;
555 }
556 }
557 };
558
559 let py_sig = &self.python_signature;
560
561 for (i, parameter) in py_sig.positional_parameters.iter().enumerate() {
562 maybe_push_comma(&mut output);
563
564 output.push_str(parameter);
565
566 if i >= py_sig.required_positional_parameters {
567 output.push('=');
568 output.push_str(&self.default_value_for_parameter(parameter));
569 }
570
571 if py_sig.positional_only_parameters > 0 && i + 1 == py_sig.positional_only_parameters {
572 output.push_str(", /")
573 }
574 }
575
576 if let Some(varargs) = &py_sig.varargs {
577 maybe_push_comma(&mut output);
578 output.push('*');
579 output.push_str(varargs);
580 } else if !py_sig.keyword_only_parameters.is_empty() {
581 maybe_push_comma(&mut output);
582 output.push('*');
583 }
584
585 for (parameter, required) in &py_sig.keyword_only_parameters {
586 maybe_push_comma(&mut output);
587 output.push_str(parameter);
588 if !required {
589 output.push('=');
590 output.push_str(&self.default_value_for_parameter(parameter));
591 }
592 }
593
594 if let Some(kwargs) = &py_sig.kwargs {
595 maybe_push_comma(&mut output);
596 output.push_str("**");
597 output.push_str(kwargs);
598 }
599
600 output.push(')');
601 output
602 }
603}
604