pyo3_stub_gen/generate/
type_alias.rs

1use std::{collections::HashSet, fmt};
2
3use crate::{
4    generate::{docstring, Import},
5    stub_type::ImportRef,
6    type_info::TypeAliasInfo,
7    TypeInfo,
8};
9
10#[derive(Debug, Clone, PartialEq)]
11pub struct TypeAliasDef {
12    pub name: &'static str,
13    pub type_: TypeInfo,
14    pub doc: &'static str,
15}
16
17impl From<&TypeAliasInfo> for TypeAliasDef {
18    fn from(info: &TypeAliasInfo) -> Self {
19        Self {
20            name: info.name,
21            type_: (info.r#type)(),
22            doc: info.doc,
23        }
24    }
25}
26
27impl Import for TypeAliasDef {
28    fn import(&self) -> HashSet<ImportRef> {
29        // Only return imports from the type itself
30        // TypeAlias will be handled conditionally by Module
31        self.type_.import.clone()
32    }
33}
34
35impl fmt::Display for TypeAliasDef {
36    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
37        write!(f, "{}: TypeAlias = {}", self.name, self.type_)
38    }
39}
40
41impl TypeAliasDef {
42    /// Format type alias with module-qualified names and syntax based on configuration
43    pub fn fmt_with_config(
44        &self,
45        target_module: &str,
46        f: &mut fmt::Formatter,
47        use_type_statement: bool,
48    ) -> fmt::Result {
49        let qualified_type = self.type_.qualified_for_module(target_module);
50
51        if use_type_statement {
52            // Python 3.12+ syntax
53            write!(f, "type {} = {}", self.name, qualified_type)?;
54        } else {
55            // Pre-3.12 syntax (default)
56            write!(f, "{}: TypeAlias = {}", self.name, qualified_type)?;
57        }
58
59        // Add docstring on next line if present
60        if !self.doc.is_empty() {
61            writeln!(f)?;
62            docstring::write_docstring(f, self.doc, "")?;
63        }
64        Ok(())
65    }
66
67    /// Existing method for backward compatibility (uses pre-3.12 syntax)
68    pub fn fmt_for_module(&self, target_module: &str, f: &mut fmt::Formatter) -> fmt::Result {
69        self.fmt_with_config(target_module, f, false)
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use std::fmt::Write;
77
78    #[test]
79    fn test_pre_312_syntax() {
80        let alias = TypeAliasDef {
81            name: "MyAlias",
82            type_: TypeInfo::builtin("int"),
83            doc: "",
84        };
85        let mut output = String::new();
86        write!(
87            &mut output,
88            "{}",
89            FormatterWrapper(&alias, "test_module", false)
90        )
91        .unwrap();
92        assert!(output.contains("MyAlias: TypeAlias = builtins.int"));
93    }
94
95    #[test]
96    fn test_312_syntax() {
97        let alias = TypeAliasDef {
98            name: "MyAlias",
99            type_: TypeInfo::builtin("int"),
100            doc: "",
101        };
102        let mut output = String::new();
103        write!(
104            &mut output,
105            "{}",
106            FormatterWrapper(&alias, "test_module", true)
107        )
108        .unwrap();
109        assert!(output.contains("type MyAlias = builtins.int"));
110        assert!(!output.contains("TypeAlias"));
111    }
112
113    // Helper struct to test formatting
114    struct FormatterWrapper<'a>(&'a TypeAliasDef, &'a str, bool);
115
116    impl<'a> fmt::Display for FormatterWrapper<'a> {
117        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
118            self.0.fmt_with_config(self.1, f, self.2)
119        }
120    }
121}
pyo3_stub_gen/generate/type_alias.rs

pyo3_stub_gen/generate/
type_alias.rs
