pyo3_stub_gen/generate/
parameters.rs

1use crate::{
2    generate::Import,
3    stub_type::ImportRef,
4    type_info::{ParameterDefault as ParameterDefaultInfo, ParameterInfo, ParameterKind},
5    TypeInfo,
6};
7use std::{collections::HashSet, fmt};
8
9/// Default value of a parameter (runtime version)
10#[derive(Debug, Clone, PartialEq)]
11pub enum ParameterDefault {
12    /// No default value
13    None,
14    /// Default value expression as a string
15    Expr(String),
16}
17
18/// A single parameter in a Python function/method signature
19///
20/// This struct represents a parameter at runtime during stub generation.
21#[derive(Debug, Clone, PartialEq)]
22pub struct Parameter {
23    /// Parameter name
24    pub name: &'static str,
25    /// Parameter kind (positional-only, keyword-only, etc.)
26    pub kind: ParameterKind,
27    /// Type information
28    pub type_info: TypeInfo,
29    /// Default value
30    pub default: ParameterDefault,
31}
32
33impl Import for Parameter {
34    fn import(&self) -> HashSet<ImportRef> {
35        self.type_info.import.clone()
36    }
37}
38
39impl From<&ParameterInfo> for Parameter {
40    fn from(info: &ParameterInfo) -> Self {
41        Self {
42            name: info.name,
43            kind: info.kind,
44            type_info: (info.type_info)(),
45            default: match &info.default {
46                ParameterDefaultInfo::None => ParameterDefault::None,
47                ParameterDefaultInfo::Expr(f) => ParameterDefault::Expr(f()),
48            },
49        }
50    }
51}
52
53impl fmt::Display for Parameter {
54    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
55        match self.kind {
56            ParameterKind::VarPositional => {
57                write!(f, "*{}: {}", self.name, self.type_info)
58            }
59            ParameterKind::VarKeyword => {
60                write!(f, "**{}: {}", self.name, self.type_info)
61            }
62            _ => {
63                write!(f, "{}: {}", self.name, self.type_info)?;
64                match &self.default {
65                    ParameterDefault::None => Ok(()),
66                    ParameterDefault::Expr(expr) => write!(f, " = {}", expr),
67                }
68            }
69        }
70    }
71}
72
73/// Container for parameters in a Python function/method signature
74///
75/// This struct organizes parameters into sections according to Python's signature syntax,
76/// ensuring proper ordering and placement of delimiters (`/` and `*`).
77#[derive(Debug, Clone, PartialEq)]
78pub struct Parameters {
79    /// Positional-only parameters (before `/`)
80    pub positional_only: Vec<Parameter>,
81    /// Positional or keyword parameters
82    pub positional_or_keyword: Vec<Parameter>,
83    /// Keyword-only parameters (after `*`)
84    pub keyword_only: Vec<Parameter>,
85    /// Variable positional parameter (`*args`)
86    pub varargs: Option<Parameter>,
87    /// Variable keyword parameter (`**kwargs`)
88    pub varkw: Option<Parameter>,
89}
90
91impl Parameters {
92    /// Create an empty Parameters container
93    pub fn new() -> Self {
94        Self {
95            positional_only: Vec::new(),
96            positional_or_keyword: Vec::new(),
97            keyword_only: Vec::new(),
98            varargs: None,
99            varkw: None,
100        }
101    }
102
103    /// Build Parameters from a slice of ParameterInfo
104    pub fn from_infos(infos: &[ParameterInfo]) -> Self {
105        let mut params = Self::new();
106
107        for info in infos {
108            let param = Parameter::from(info);
109            match param.kind {
110                ParameterKind::PositionalOnly => params.positional_only.push(param),
111                ParameterKind::PositionalOrKeyword => params.positional_or_keyword.push(param),
112                ParameterKind::KeywordOnly => params.keyword_only.push(param),
113                ParameterKind::VarPositional => params.varargs = Some(param),
114                ParameterKind::VarKeyword => params.varkw = Some(param),
115            }
116        }
117
118        params
119    }
120
121    /// Iterate over all parameters in signature order
122    pub fn iter_entries(&self) -> impl Iterator<Item = &Parameter> {
123        self.positional_only
124            .iter()
125            .chain(self.positional_or_keyword.iter())
126            .chain(self.varargs.iter())
127            .chain(self.keyword_only.iter())
128            .chain(self.varkw.iter())
129    }
130
131    /// Check if there are no parameters at all
132    pub fn is_empty(&self) -> bool {
133        self.positional_only.is_empty()
134            && self.positional_or_keyword.is_empty()
135            && self.keyword_only.is_empty()
136            && self.varargs.is_none()
137            && self.varkw.is_none()
138    }
139}
140
141impl Default for Parameters {
142    fn default() -> Self {
143        Self::new()
144    }
145}
146
147impl Import for Parameters {
148    fn import(&self) -> HashSet<ImportRef> {
149        self.iter_entries().flat_map(|p| p.import()).collect()
150    }
151}
152
153impl fmt::Display for Parameters {
154    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
155        let mut parts = Vec::new();
156
157        // Positional-only parameters
158        for param in &self.positional_only {
159            parts.push(param.to_string());
160        }
161
162        // Insert `/` delimiter if there are positional-only parameters
163        if !self.positional_only.is_empty() {
164            parts.push("/".to_string());
165        }
166
167        // Positional or keyword parameters
168        for param in &self.positional_or_keyword {
169            parts.push(param.to_string());
170        }
171
172        // Variable positional parameter or bare `*` for keyword-only
173        if let Some(varargs) = &self.varargs {
174            parts.push(varargs.to_string());
175        } else if !self.keyword_only.is_empty() {
176            // Need bare `*` to indicate keyword-only parameters follow
177            parts.push("*".to_string());
178        }
179
180        // Keyword-only parameters
181        for param in &self.keyword_only {
182            parts.push(param.to_string());
183        }
184
185        // Variable keyword parameter
186        if let Some(varkw) = &self.varkw {
187            parts.push(varkw.to_string());
188        }
189
190        write!(f, "{}", parts.join(", "))
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197
198    #[test]
199    fn test_positional_only() {
200        let params = Parameters {
201            positional_only: vec![
202                Parameter {
203                    name: "x",
204                    kind: ParameterKind::PositionalOnly,
205                    type_info: TypeInfo::builtin("int"),
206                    default: ParameterDefault::None,
207                },
208                Parameter {
209                    name: "y",
210                    kind: ParameterKind::PositionalOnly,
211                    type_info: TypeInfo::builtin("int"),
212                    default: ParameterDefault::None,
213                },
214            ],
215            ..Default::default()
216        };
217
218        assert_eq!(params.to_string(), "x: builtins.int, y: builtins.int, /");
219    }
220
221    #[test]
222    fn test_keyword_only() {
223        let params = Parameters {
224            keyword_only: vec![Parameter {
225                name: "kw",
226                kind: ParameterKind::KeywordOnly,
227                type_info: TypeInfo::builtin("str"),
228                default: ParameterDefault::Expr("None".to_string()),
229            }],
230            ..Default::default()
231        };
232
233        assert_eq!(params.to_string(), "*, kw: builtins.str = None");
234    }
235
236    #[test]
237    fn test_mixed_with_defaults() {
238        let params = Parameters {
239            positional_only: vec![Parameter {
240                name: "token",
241                kind: ParameterKind::PositionalOnly,
242                type_info: TypeInfo::builtin("str"),
243                default: ParameterDefault::None,
244            }],
245            keyword_only: vec![
246                Parameter {
247                    name: "retries",
248                    kind: ParameterKind::KeywordOnly,
249                    type_info: TypeInfo::builtin("int"),
250                    default: ParameterDefault::Expr("3".to_string()),
251                },
252                Parameter {
253                    name: "timeout",
254                    kind: ParameterKind::KeywordOnly,
255                    type_info: TypeInfo::builtin("float"),
256                    default: ParameterDefault::None,
257                },
258            ],
259            ..Default::default()
260        };
261
262        assert_eq!(
263            params.to_string(),
264            "token: builtins.str, /, *, retries: builtins.int = 3, timeout: builtins.float"
265        );
266    }
267
268    #[test]
269    fn test_varargs_kwargs() {
270        let params = Parameters {
271            varargs: Some(Parameter {
272                name: "args",
273                kind: ParameterKind::VarPositional,
274                type_info: TypeInfo::builtin("str"),
275                default: ParameterDefault::None,
276            }),
277            varkw: Some(Parameter {
278                name: "kwargs",
279                kind: ParameterKind::VarKeyword,
280                type_info: TypeInfo::any(),
281                default: ParameterDefault::None,
282            }),
283            ..Default::default()
284        };
285
286        assert_eq!(
287            params.to_string(),
288            "*args: builtins.str, **kwargs: typing.Any"
289        );
290    }
291}
pyo3_stub_gen/generate/parameters.rs

pyo3_stub_gen/generate/
parameters.rs
