1use std::collections::HashMap;
2
3use proc_macro2::TokenStream as TokenStream2;
4use quote::{quote, ToTokens, TokenStreamExt};
5use syn::{
6 parenthesized,
7 parse::{Parse, ParseStream},
8 punctuated::Punctuated,
9 token, Expr, Ident, Result, Token, Type,
10};
11
12use crate::gen_stub::remove_lifetime;
13
14use super::ArgInfo;
15
16#[derive(Debug, Clone, PartialEq)]
17enum SignatureArg {
18 Ident(Ident),
19 Assign(Ident, Token![=], Expr),
20 Star(Token![*]),
21 Args(Token![*], Ident),
22 Keywords(Token![*], Token![*], Ident),
23}
24
25impl Parse for SignatureArg {
26 fn parse(input: ParseStream) -> Result<Self> {
27 if input.peek(Token![*]) {
28 let star: Star = input.parse()?;
29 if input.peek(Token![*]) {
30 Ok(SignatureArg::Keywords(star, input.parse()?, input.parse()?))
31 } else if input.peek(token:Ident) {
32 Ok(SignatureArg::Args(star, input.parse()?))
33 } else {
34 Ok(SignatureArg::Star(star))
35 }
36 } else if input.peek(token:Ident) {
37 let ident: Ident = Ident::parse(input)?;
38 if input.peek(Token![=]) {
39 Ok(SignatureArg::Assign(ident, input.parse()?, input.parse()?))
40 } else {
41 Ok(SignatureArg::Ident(ident))
42 }
43 } else {
44 dbg!(input);
45 todo!()
46 }
47 }
48}
49
50#[derive(Debug, Clone, PartialEq)]
51pub struct Signature {
52 paren: token::Paren,
53 args: Punctuated<SignatureArg, Token![,]>,
54}
55
56impl Parse for Signature {
57 fn parse(input: ParseStream) -> Result<Self> {
58 let content: ParseBuffer<'_>;
59 let paren: Paren = parenthesized!(content in input);
60 let args: Punctuated = content.parse_terminated(parser:SignatureArg::parse, separator:Token![,])?;
61 Ok(Self { paren, args })
62 }
63}
64
65pub struct ArgsWithSignature<'a> {
66 pub args: &'a Vec<ArgInfo>,
67 pub sig: &'a Option<Signature>,
68}
69
70impl ToTokens for ArgsWithSignature<'_> {
71 fn to_tokens(&self, tokens: &mut TokenStream2) {
72 let arg_infos: Vec<TokenStream2> = if let Some(sig) = self.sig {
73 // record all Type information from rust's args
74 let args_map: HashMap<String, Type> = self
75 .args
76 .iter()
77 .map(|arg| {
78 let mut ty = arg.r#type.clone();
79 remove_lifetime(&mut ty);
80 (arg.name.clone(), ty)
81 })
82 .collect();
83 sig.args.iter().map(|sig_arg| match sig_arg {
84 SignatureArg::Ident(ident) => {
85 let name = ident.to_string();
86 let ty = args_map.get(&name).unwrap();
87 quote! {
88 ::pyo3_stub_gen::type_info::ArgInfo {
89 name: #name,
90 r#type: <#ty as ::pyo3_stub_gen::PyStubType>::type_input,
91 signature: Some(pyo3_stub_gen::type_info::SignatureArg::Ident),
92 }
93 }
94 }
95 SignatureArg::Assign(ident, _eq, value) => {
96 let name = ident.to_string();
97 let ty = args_map.get(&name).unwrap();
98 quote! {
99 ::pyo3_stub_gen::type_info::ArgInfo {
100 name: #name,
101 r#type: <#ty as ::pyo3_stub_gen::PyStubType>::type_input,
102 signature: Some(pyo3_stub_gen::type_info::SignatureArg::Assign{
103 default: {
104 static DEFAULT: std::sync::LazyLock<String> = std::sync::LazyLock::new(|| {
105 ::pyo3::prepare_freethreaded_python();
106 ::pyo3::Python::with_gil(|py| -> String {
107 let v: #ty = #value;
108 if let Ok(py_obj) = <#ty as ::pyo3::IntoPyObjectExt>::into_bound_py_any(v, py) {
109 ::pyo3_stub_gen::util::fmt_py_obj(&py_obj)
110 } else {
111 "...".to_owned()
112 }
113 })
114 });
115 &DEFAULT
116 }
117 }),
118 }
119 }
120 },
121 SignatureArg::Star(_) => quote! {
122 ::pyo3_stub_gen::type_info::ArgInfo {
123 name: "",
124 r#type: <() as ::pyo3_stub_gen::PyStubType>::type_input,
125 signature: Some(pyo3_stub_gen::type_info::SignatureArg::Star),
126 }
127 },
128 SignatureArg::Args(_, ident) => {
129 let name = ident.to_string();
130 let ty = args_map.get(&name).unwrap();
131 quote! {
132 ::pyo3_stub_gen::type_info::ArgInfo {
133 name: #name,
134 r#type: <#ty as ::pyo3_stub_gen::PyStubType>::type_input,
135 signature: Some(pyo3_stub_gen::type_info::SignatureArg::Args),
136 }
137 }
138 },
139 SignatureArg::Keywords(_, _, ident) => {
140 let name = ident.to_string();
141 let ty = args_map.get(&name).unwrap();
142 quote! {
143 ::pyo3_stub_gen::type_info::ArgInfo {
144 name: #name,
145 r#type: <#ty as ::pyo3_stub_gen::PyStubType>::type_input,
146 signature: Some(pyo3_stub_gen::type_info::SignatureArg::Keywords),
147 }
148 }
149 }
150 }).collect()
151 } else {
152 self.args
153 .iter()
154 .map(|arg| {
155 let mut ty = arg.r#type.clone();
156 remove_lifetime(&mut ty);
157 let name = &arg.name;
158 quote! {
159 ::pyo3_stub_gen::type_info::ArgInfo {
160 name: #name,
161 r#type: <#ty as ::pyo3_stub_gen::PyStubType>::type_input,
162 signature: None,
163 }
164 }
165 })
166 .collect()
167 };
168 tokens.append_all(quote! { &[ #(#arg_infos),* ] });
169 }
170}
171
172impl Signature {
173 pub fn overriding_operator(sig: &syn::Signature) -> Option<Self> {
174 if sig.ident == "__pow__" {
175 return Some(syn::parse_str("(exponent, modulo=None)").unwrap());
176 }
177 if sig.ident == "__rpow__" {
178 return Some(syn::parse_str("(base, modulo=None)").unwrap());
179 }
180 None
181 }
182}
183