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
194impl Parameters {
195    /// Format parameters with module-qualified type names
196    ///
197    /// This method uses the target module context to qualify type identifiers
198    /// within compound type expressions based on their source modules.
199    pub fn fmt_for_module(&self, target_module: &str) -> String {
200        let mut parts = Vec::new();
201
202        // Positional-only parameters
203        for param in &self.positional_only {
204            parts.push(Self::format_param(param, target_module));
205        }
206
207        // Insert `/` delimiter if there are positional-only parameters
208        if !self.positional_only.is_empty() {
209            parts.push("/".to_string());
210        }
211
212        // Positional or keyword parameters
213        for param in &self.positional_or_keyword {
214            parts.push(Self::format_param(param, target_module));
215        }
216
217        // Variable positional parameter or bare `*` for keyword-only
218        if let Some(varargs) = &self.varargs {
219            parts.push(Self::format_param(varargs, target_module));
220        } else if !self.keyword_only.is_empty() {
221            // Need bare `*` to indicate keyword-only parameters follow
222            parts.push("*".to_string());
223        }
224
225        // Keyword-only parameters
226        for param in &self.keyword_only {
227            parts.push(Self::format_param(param, target_module));
228        }
229
230        // Variable keyword parameter
231        if let Some(varkw) = &self.varkw {
232            parts.push(Self::format_param(varkw, target_module));
233        }
234
235        parts.join(", ")
236    }
237
238    /// Format a single parameter with qualified type names
239    fn format_param(param: &Parameter, target_module: &str) -> String {
240        let qualified_type = param.type_info.qualified_for_module(target_module);
241        match param.kind {
242            ParameterKind::VarPositional => format!("*{}: {}", param.name, qualified_type),
243            ParameterKind::VarKeyword => format!("**{}: {}", param.name, qualified_type),
244            _ => {
245                let base = format!("{}: {}", param.name, qualified_type);
246                match &param.default {
247                    ParameterDefault::None => base,
248                    ParameterDefault::Expr(expr) => format!("{} = {}", base, expr),
249                }
250            }
251        }
252    }
253}
254
255#[cfg(test)]
256mod tests {
257    use super::*;
258
259    #[test]
260    fn test_positional_only() {
261        let params = Parameters {
262            positional_only: vec![
263                Parameter {
264                    name: "x",
265                    kind: ParameterKind::PositionalOnly,
266                    type_info: TypeInfo::builtin("int"),
267                    default: ParameterDefault::None,
268                },
269                Parameter {
270                    name: "y",
271                    kind: ParameterKind::PositionalOnly,
272                    type_info: TypeInfo::builtin("int"),
273                    default: ParameterDefault::None,
274                },
275            ],
276            ..Default::default()
277        };
278
279        assert_eq!(params.to_string(), "x: builtins.int, y: builtins.int, /");
280    }
281
282    #[test]
283    fn test_keyword_only() {
284        let params = Parameters {
285            keyword_only: vec![Parameter {
286                name: "kw",
287                kind: ParameterKind::KeywordOnly,
288                type_info: TypeInfo::builtin("str"),
289                default: ParameterDefault::Expr("None".to_string()),
290            }],
291            ..Default::default()
292        };
293
294        assert_eq!(params.to_string(), "*, kw: builtins.str = None");
295    }
296
297    #[test]
298    fn test_mixed_with_defaults() {
299        let params = Parameters {
300            positional_only: vec![Parameter {
301                name: "token",
302                kind: ParameterKind::PositionalOnly,
303                type_info: TypeInfo::builtin("str"),
304                default: ParameterDefault::None,
305            }],
306            keyword_only: vec![
307                Parameter {
308                    name: "retries",
309                    kind: ParameterKind::KeywordOnly,
310                    type_info: TypeInfo::builtin("int"),
311                    default: ParameterDefault::Expr("3".to_string()),
312                },
313                Parameter {
314                    name: "timeout",
315                    kind: ParameterKind::KeywordOnly,
316                    type_info: TypeInfo::builtin("float"),
317                    default: ParameterDefault::None,
318                },
319            ],
320            ..Default::default()
321        };
322
323        assert_eq!(
324            params.to_string(),
325            "token: builtins.str, /, *, retries: builtins.int = 3, timeout: builtins.float"
326        );
327    }
328
329    #[test]
330    fn test_varargs_kwargs() {
331        let params = Parameters {
332            varargs: Some(Parameter {
333                name: "args",
334                kind: ParameterKind::VarPositional,
335                type_info: TypeInfo::builtin("str"),
336                default: ParameterDefault::None,
337            }),
338            varkw: Some(Parameter {
339                name: "kwargs",
340                kind: ParameterKind::VarKeyword,
341                type_info: TypeInfo::any(),
342                default: ParameterDefault::None,
343            }),
344            ..Default::default()
345        };
346
347        assert_eq!(
348            params.to_string(),
349            "*args: builtins.str, **kwargs: typing.Any"
350        );
351    }
352}
pyo3_stub_gen/generate/parameters.rs

pyo3_stub_gen/generate/
parameters.rs
