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