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                Token::NumericLiteral(num) => result.push_str(num),
208            }
209        }
210        result
211    }
212}
213
214/// Renderer for type expressions
215pub struct TypeRenderer<'a> {
216    link_resolver: &'a LinkResolver<'a>,
217    current_module: &'a str,
218}
219
220impl<'a> TypeRenderer<'a> {
221    pub fn new(link_resolver: &'a LinkResolver<'a>, current_module: &'a str) -> Self {
222        Self {
223            link_resolver,
224            current_module,
225        }
226    }
227
228    /// Render a type expression
229    /// 1. Check if this is a type alias - preserve name, don't expand
230    /// 2. Strip module prefixes from display text
231    /// 3. Resolve link target using LinkResolver
232    /// 4. Recursively handle generic parameters
233    pub fn render_type(&self, type_info: &TypeInfo) -> DocTypeExpr {
234        // Parse the type expression into a structure
235        let type_structure = TypeStructure::parse(&type_info.name);
236
237        // Recursively render the structure
238        self.render_type_structure(&type_structure, type_info)
239    }
240
241    /// Recursively render a TypeStructure into a DocTypeExpr
242    fn render_type_structure(
243        &self,
244        structure: &TypeStructure,
245        type_info: &TypeInfo,
246    ) -> DocTypeExpr {
247        match structure {
248            TypeStructure::Simple(name) => {
249                // Simple type - try to create a link
250                let display = self.strip_module_prefix(name);
251                let link_target = self.try_create_link_for_name(name, type_info);
252
253                DocTypeExpr {
254                    display,
255                    link_target,
256                    children: Vec::new(),
257                }
258            }
259
260            TypeStructure::Generic { base, args } => {
261                // Generic type - render base and recursively render args
262                let base_display = self.strip_module_prefix(base);
263                let base_link = self.try_create_link_for_name(base, type_info);
264
265                // Recursively render all arguments
266                let children: Vec<DocTypeExpr> = args
267                    .iter()
268                    .map(|arg| self.render_type_structure(arg, type_info))
269                    .collect();
270
271                // Build display string with brackets
272                let args_display: Vec<String> =
273                    children.iter().map(|child| child.display.clone()).collect();
274                let display = format!("{}[{}]", base_display, args_display.join(", "));
275
276                DocTypeExpr {
277                    display,
278                    link_target: base_link,
279                    children,
280                }
281            }
282
283            TypeStructure::Union(variants) => {
284                // Union type - recursively render all variants
285                let children: Vec<DocTypeExpr> = variants
286                    .iter()
287                    .map(|variant| self.render_type_structure(variant, type_info))
288                    .collect();
289
290                // Build display string with pipe operators
291                let variants_display: Vec<String> =
292                    children.iter().map(|child| child.display.clone()).collect();
293                let display = variants_display.join(" | ");
294
295                DocTypeExpr {
296                    display,
297                    link_target: None, // Union itself has no link
298                    children,
299                }
300            }
301        }
302    }
303
304    /// Try to create a link target for a simple type name
305    /// Uses type_refs to look up module information
306    fn try_create_link_for_name(&self, name: &str, type_info: &TypeInfo) -> Option<LinkTarget> {
307        // Extract bare identifier from potentially qualified name
308        // e.g., "main_mod.A" -> "A"
309        let bare_name = name.split('.').next_back().unwrap_or(name);
310
311        // Look up in type_refs
312        let type_ref = type_info.type_refs.get(bare_name)?;
313
314        // Get module path from TypeIdentifierRef
315        let module_path = type_ref.module.get()?;
316
317        // Create FQN (fully qualified name)
318        let fqn = format!("{}.{}", module_path, bare_name);
319
320        // Use export_map to find the correct doc module
321        let doc_module = self
322            .link_resolver
323            .export_map()
324            .get(&fqn)
325            .cloned()
326            .unwrap_or_else(|| self.current_module.to_string());
327
328        Some(LinkTarget {
329            fqn,
330            doc_module,
331            kind: ItemKind::Class, // Assume class for now
332            attribute: None,
333        })
334    }
335
336    /// Strip module prefixes from type names
337    /// Remove "typing.", "builtins.", "package.submod."
338    /// Keep only bare names: "Optional[ClassA]" not "typing.Optional[sub_mod.ClassA]"
339    fn strip_module_prefix(&self, type_name: &str) -> String {
340        prefix_stripper::strip_type_prefixes(type_name, self.current_module)
341    }
342}
pyo3_stub_gen/docgen/types.rs

pyo3_stub_gen/docgen/
types.rs
