pyo3_macros_backend/pyfunction/
signature.rs

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

pyo3_macros_backend/pyfunction/
signature.rs
