pyo3_macros_backend/pyfunction/
signature.rs

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;
25        let paren_token = syn::parenthesized!(content in input);
26
27        let items = content.parse_terminated(SignatureItem::parse, 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| 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 = input.lookahead1();
80        if lookahead.peek(Token![*]) {
81            if input.peek2(Token![*]) {
82                input.parse().map(SignatureItem::Kwargs)
83            } else {
84                let sep = 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(SignatureItem::PosargsSep)
96        } else {
97            input.parse().map(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) => arg.to_tokens(tokens),
106            SignatureItem::Varargs(varargs) => varargs.to_tokens(tokens),
107            SignatureItem::VarargsSep(sep) => sep.to_tokens(tokens),
108            SignatureItem::Kwargs(kwargs) => kwargs.to_tokens(tokens),
109            SignatureItem::PosargsSep(sep) => 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, default)) = &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}
pyo3_macros_backend/pyfunction/signature.rs

pyo3_macros_backend/pyfunction/
signature.rs
