pyo3_stub_gen/docgen/
types.rs

1//! Type expression rendering for documentation
2
3use crate::docgen::ir::{DocTypeExpr, ItemKind, LinkTarget};
4use crate::docgen::link::LinkResolver;
5use crate::docgen::util::prefix_stripper;
6use crate::generate::qualifier::{tokenize, Token};
7use crate::TypeInfo;
8
9/// Type alias definition placeholder
10/// TODO: This will be used for type alias expansion in future versions
11pub struct TypeAliasDef;
12
13/// Parsed structure of a type expression
14#[derive(Debug, Clone, PartialEq)]
15enum TypeStructure {
16    /// Simple type: "int", "ClassA"
17    Simple(String),
18    /// Generic type: "Optional[A]", "Dict[str, int]"
19    Generic {
20        base: String,
21        args: Vec<TypeStructure>,
22    },
23    /// Union type: "A | B | C"
24    Union(Vec<TypeStructure>),
25}
26
27impl TypeStructure {
28    /// Parse a type expression string into a TypeStructure
29    fn parse(expr: &str) -> Self {
30        let tokens = tokenize(expr);
31        Self::parse_tokens(&tokens)
32    }
33
34    /// Parse a slice of tokens into a TypeStructure
35    fn parse_tokens(tokens: &[Token]) -> Self {
36        // Filter out whitespace for easier parsing
37        let filtered_tokens: Vec<Token> = tokens
38            .iter()
39            .filter(|t| !matches!(t, Token::Whitespace(_)))
40            .cloned()
41            .collect();
42
43        if filtered_tokens.is_empty() {
44            return TypeStructure::Simple(String::new());
45        }
46
47        // Try to parse as union first (lowest precedence)
48        if let Some(union_parts) = Self::try_parse_union(&filtered_tokens) {
49            let variants: Vec<TypeStructure> = union_parts
50                .into_iter()
51                .map(|part| Self::parse_tokens(&part))
52                .collect();
53            return TypeStructure::Union(variants);
54        }
55
56        // Try to parse as generic
57        if let Some((base_tokens, arg_token_lists)) = Self::try_parse_generic(&filtered_tokens) {
58            let base = Self::tokens_to_string(&base_tokens);
59            let args: Vec<TypeStructure> = arg_token_lists
60                .into_iter()
61                .map(|arg_tokens| Self::parse_tokens(&arg_tokens))
62                .collect();
63            return TypeStructure::Generic { base, args };
64        }
65
66        // Simple type
67        TypeStructure::Simple(Self::tokens_to_string(&filtered_tokens))
68    }
69
70    /// Try to parse tokens as a union (A | B | C)
71    /// Returns Some(vec![tokens_for_A, tokens_for_B, tokens_for_C]) if successful
72    fn try_parse_union(tokens: &[Token]) -> Option<Vec<Vec<Token>>> {
73        let mut parts = Vec::new();
74        let mut current_part = Vec::new();
75        let mut depth = 0;
76
77        for token in tokens {
78            match token {
79                Token::OpenBracket(_) => {
80                    depth += 1;
81                    current_part.push(token.clone());
82                }
83                Token::CloseBracket(_) => {
84                    depth -= 1;
85                    current_part.push(token.clone());
86                }
87                Token::Pipe if depth == 0 => {
88                    if !current_part.is_empty() {
89                        parts.push(current_part);
90                        current_part = Vec::new();
91                    }
92                }
93                _ => {
94                    current_part.push(token.clone());
95                }
96            }
97        }
98
99        if !current_part.is_empty() {
100            parts.push(current_part);
101        }
102
103        // Only return Some if we actually found pipes (i.e., more than one part)
104        if parts.len() > 1 {
105            Some(parts)
106        } else {
107            None
108        }
109    }
110
111    /// Try to parse tokens as a generic type
112    /// Returns Some((base_tokens, vec![arg1_tokens, arg2_tokens, ...])) if successful
113    fn try_parse_generic(tokens: &[Token]) -> Option<(Vec<Token>, Vec<Vec<Token>>)> {
114        // Find the first opening bracket
115        let bracket_pos = tokens
116            .iter()
117            .position(|t| matches!(t, Token::OpenBracket('[')))?;
118
119        // Everything before the bracket is the base
120        let base_tokens = tokens[..bracket_pos].to_vec();
121
122        // Find matching closing bracket
123        let mut depth = 0;
124        let mut close_pos = None;
125        for (i, token) in tokens[bracket_pos..].iter().enumerate() {
126            match token {
127                Token::OpenBracket(_) => depth += 1,
128                Token::CloseBracket(']') => {
129                    depth -= 1;
130                    if depth == 0 {
131                        close_pos = Some(bracket_pos + i);
132                        break;
133                    }
134                }
135                _ => {}
136            }
137        }
138
139        let close_pos = close_pos?;
140
141        // Extract the contents inside the brackets
142        let inner_tokens = &tokens[bracket_pos + 1..close_pos];
143
144        // Split by commas at depth 0
145        let arg_token_lists = Self::split_by_comma_at_depth_zero(inner_tokens);
146
147        Some((base_tokens, arg_token_lists))
148    }
149
150    /// Split tokens by commas at depth 0
151    fn split_by_comma_at_depth_zero(tokens: &[Token]) -> Vec<Vec<Token>> {
152        let mut parts = Vec::new();
153        let mut current_part = Vec::new();
154        let mut depth = 0;
155
156        for token in tokens {
157            match token {
158                Token::OpenBracket(_) => {
159                    depth += 1;
160                    current_part.push(token.clone());
161                }
162                Token::CloseBracket(_) => {
163                    depth -= 1;
164                    current_part.push(token.clone());
165                }
166                Token::Comma if depth == 0 => {
167                    if !current_part.is_empty() {
168                        parts.push(current_part);
169                        current_part = Vec::new();
170                    }
171                }
172                _ => {
173                    current_part.push(token.clone());
174                }
175            }
176        }
177
178        if !current_part.is_empty() {
179            parts.push(current_part);
180        }
181
182        if parts.is_empty() {
183            vec![tokens.to_vec()]
184        } else {
185            parts
186        }
187    }
188
189    /// Convert tokens to a string representation
190    fn tokens_to_string(tokens: &[Token]) -> String {
191        let mut result = String::new();
192        for token in tokens {
193            match token {
194                Token::Identifier(s) => result.push_str(s),
195                Token::DottedPath(parts) => result.push_str(&parts.join(".")),
196                Token::OpenBracket(ch) => result.push(*ch),
197                Token::CloseBracket(ch) => result.push(*ch),
198                Token::Comma => result.push(','),
199                Token::Pipe => result.push_str(" | "),
200                Token::Ellipsis => result.push_str("..."),
201                Token::StringLiteral(s) => {
202                    result.push('"');
203                    result.push_str(s);
204                    result.push('"');
205                }
206                Token::Whitespace(ws) => result.push_str(ws),
207            }
208        }
209        result
210    }
211}
212
213/// Renderer for type expressions
214pub struct TypeRenderer<'a> {
215    link_resolver: &'a LinkResolver<'a>,
216    current_module: &'a str,
217}
218
219impl<'a> TypeRenderer<'a> {
220    pub fn new(link_resolver: &'a LinkResolver<'a>, current_module: &'a str) -> Self {
221        Self {
222            link_resolver,
223            current_module,
224        }
225    }
226
227    /// Render a type expression
228    /// 1. Check if this is a type alias - preserve name, don't expand
229    /// 2. Strip module prefixes from display text
230    /// 3. Resolve link target using LinkResolver
231    /// 4. Recursively handle generic parameters
232    pub fn render_type(&self, type_info: &TypeInfo) -> DocTypeExpr {
233        // Parse the type expression into a structure
234        let type_structure = TypeStructure::parse(&type_info.name);
235
236        // Recursively render the structure
237        self.render_type_structure(&type_structure, type_info)
238    }
239
240    /// Recursively render a TypeStructure into a DocTypeExpr
241    fn render_type_structure(
242        &self,
243        structure: &TypeStructure,
244        type_info: &TypeInfo,
245    ) -> DocTypeExpr {
246        match structure {
247            TypeStructure::Simple(name) => {
248                // Simple type - try to create a link
249                let display = self.strip_module_prefix(name);
250                let link_target = self.try_create_link_for_name(name, type_info);
251
252                DocTypeExpr {
253                    display,
254                    link_target,
255                    children: Vec::new(),
256                }
257            }
258
259            TypeStructure::Generic { base, args } => {
260                // Generic type - render base and recursively render args
261                let base_display = self.strip_module_prefix(base);
262                let base_link = self.try_create_link_for_name(base, type_info);
263
264                // Recursively render all arguments
265                let children: Vec<DocTypeExpr> = args
266                    .iter()
267                    .map(|arg| self.render_type_structure(arg, type_info))
268                    .collect();
269
270                // Build display string with brackets
271                let args_display: Vec<String> =
272                    children.iter().map(|child| child.display.clone()).collect();
273                let display = format!("{}[{}]", base_display, args_display.join(", "));
274
275                DocTypeExpr {
276                    display,
277                    link_target: base_link,
278                    children,
279                }
280            }
281
282            TypeStructure::Union(variants) => {
283                // Union type - recursively render all variants
284                let children: Vec<DocTypeExpr> = variants
285                    .iter()
286                    .map(|variant| self.render_type_structure(variant, type_info))
287                    .collect();
288
289                // Build display string with pipe operators
290                let variants_display: Vec<String> =
291                    children.iter().map(|child| child.display.clone()).collect();
292                let display = variants_display.join(" | ");
293
294                DocTypeExpr {
295                    display,
296                    link_target: None, // Union itself has no link
297                    children,
298                }
299            }
300        }
301    }
302
303    /// Try to create a link target for a simple type name
304    /// Uses type_refs to look up module information
305    fn try_create_link_for_name(&self, name: &str, type_info: &TypeInfo) -> Option<LinkTarget> {
306        // Extract bare identifier from potentially qualified name
307        // e.g., "main_mod.A" -> "A"
308        let bare_name = name.split('.').next_back().unwrap_or(name);
309
310        // Look up in type_refs
311        let type_ref = type_info.type_refs.get(bare_name)?;
312
313        // Get module path from TypeIdentifierRef
314        let module_path = type_ref.module.get()?;
315
316        // Create FQN (fully qualified name)
317        let fqn = format!("{}.{}", module_path, bare_name);
318
319        // Use export_map to find the correct doc module
320        let doc_module = self
321            .link_resolver
322            .export_map()
323            .get(&fqn)
324            .cloned()
325            .unwrap_or_else(|| self.current_module.to_string());
326
327        Some(LinkTarget {
328            fqn,
329            doc_module,
330            kind: ItemKind::Class, // Assume class for now
331            attribute: None,
332        })
333    }
334
335    /// Strip module prefixes from type names
336    /// Remove "typing.", "builtins.", "package.submod."
337    /// Keep only bare names: "Optional[ClassA]" not "typing.Optional[sub_mod.ClassA]"
338    fn strip_module_prefix(&self, type_name: &str) -> String {
339        prefix_stripper::strip_type_prefixes(type_name, self.current_module)
340    }
341}
pyo3_stub_gen/docgen/types.rs

pyo3_stub_gen/docgen/
types.rs
