pyo3_stub_gen/docgen/
builder.rs

1//! Builder for converting StubInfo to DocPackage
2
3use crate::docgen::{
4    export::ExportResolver,
5    ir::{
6        DeprecatedInfo, DocAttribute, DocClass, DocFunction, DocItem, DocModule, DocPackage,
7        DocParameter, DocSignature, DocSubmodule, DocTypeAlias, DocTypeExpr, DocVariable,
8    },
9    types::TypeRenderer,
10    util::{is_hidden_module, prefix_stripper},
11};
12use crate::generate::StubInfo;
13use crate::Result;
14use std::collections::BTreeMap;
15use std::path::PathBuf;
16
17/// Helper to check if item already exists in the list
18fn matches_item_name(item: &DocItem, name: &str) -> bool {
19    match item {
20        DocItem::Function(f) => f.name == name,
21        DocItem::Class(c) => c.name == name,
22        DocItem::TypeAlias(t) => t.name == name,
23        DocItem::Variable(v) => v.name == name,
24        DocItem::Module(m) => m.name == name,
25    }
26}
27
28/// Context for building documentation items, containing shared rendering components
29struct DocBuildContext<'a> {
30    link_resolver: crate::docgen::link::LinkResolver<'a>,
31    module: &'a str,
32}
33
34impl<'a> DocBuildContext<'a> {
35    /// Get a type renderer for this context
36    fn type_renderer(&self) -> TypeRenderer<'_> {
37        TypeRenderer::new(&self.link_resolver, self.module)
38    }
39
40    /// Get a default value parser for this context
41    fn default_parser(&self) -> crate::docgen::default_parser::DefaultValueParser<'_> {
42        crate::docgen::default_parser::DefaultValueParser::new(&self.link_resolver, self.module)
43    }
44}
45
46/// Builder for converting StubInfo to DocPackage
47pub struct DocPackageBuilder<'a> {
48    stub_info: &'a StubInfo,
49    export_resolver: ExportResolver<'a>,
50    export_map: BTreeMap<String, String>,
51    default_module_name: String,
52}
53
54impl<'a> DocPackageBuilder<'a> {
55    pub fn new(stub_info: &'a StubInfo) -> Self {
56        let export_resolver = ExportResolver::new(&stub_info.modules);
57        let export_map = export_resolver.build_export_map();
58
59        // Get the default module name from the first module (they all share the same default_module_name)
60        let default_module_name = stub_info
61            .modules
62            .values()
63            .next()
64            .map(|m| m.default_module_name.clone())
65            .unwrap_or_default();
66
67        Self {
68            stub_info,
69            export_resolver,
70            export_map,
71            default_module_name,
72        }
73    }
74
75    /// Create a build context for a specific module
76    fn create_context<'b>(&'b self, module: &'b str) -> DocBuildContext<'b> {
77        let link_resolver = crate::docgen::link::LinkResolver::new(&self.export_map);
78        DocBuildContext {
79            link_resolver,
80            module,
81        }
82    }
83
84    pub fn build(self) -> Result<DocPackage> {
85        let mut modules = BTreeMap::new();
86
87        for (module_name, module) in &self.stub_info.modules {
88            // Skip modules with any component starting with '_'
89            if is_hidden_module(module_name) {
90                continue;
91            }
92
93            let doc_module = self.build_module(module_name, module)?;
94            modules.insert(module_name.clone(), doc_module);
95        }
96
97        let export_map = self.export_map;
98
99        // Convert config to use relative POSIX path for JSON serialization
100        let mut json_config = self.stub_info.config.doc_gen.clone().unwrap_or_default();
101        if let Some(pyproject_dir) = &self.stub_info.pyproject_dir {
102            let relative_posix = json_config.to_relative_posix_path(pyproject_dir);
103            json_config.output_dir = PathBuf::from(relative_posix);
104        }
105
106        Ok(DocPackage {
107            name: self.default_module_name.clone(),
108            modules,
109            export_map,
110            config: json_config,
111        })
112    }
113
114    fn build_module(&self, name: &str, module: &crate::generate::Module) -> Result<DocModule> {
115        let exports = self.export_resolver.resolve_exports(module);
116        let mut items = Vec::new();
117
118        // Process functions - handle overloads (Requirement #1)
119        for (func_name, func_defs) in &module.function {
120            if exports.contains(*func_name) {
121                items.push(self.build_function(name, func_defs)?);
122            }
123        }
124
125        // Process type aliases (Requirement #2)
126        for (alias_name, alias_def) in &module.type_aliases {
127            if exports.contains(*alias_name) {
128                items.push(self.build_type_alias(name, alias_def)?);
129            }
130        }
131
132        // Process classes (sorted by name for deterministic output)
133        let mut classes: Vec<_> = module.class.values().collect();
134        classes.sort_by_key(|c| c.name);
135        for class_def in classes {
136            if exports.contains(class_def.name) {
137                items.push(self.build_class(name, class_def)?);
138            }
139        }
140
141        // Process enums (sorted by name for deterministic output)
142        let mut enums: Vec<_> = module.enum_.values().collect();
143        enums.sort_by_key(|e| e.name);
144        for enum_def in enums {
145            if exports.contains(enum_def.name) {
146                items.push(self.build_enum_as_class(name, enum_def)?);
147            }
148        }
149
150        // Process variables
151        for (var_name, var_def) in &module.variables {
152            if exports.contains(*var_name) {
153                items.push(self.build_variable(name, var_def)?);
154            }
155        }
156
157        // Process module re-exports (from reexport_module_members!)
158        for re_export in &module.module_re_exports {
159            if let Some(source_module) = self.stub_info.modules.get(&re_export.source_module) {
160                for item_name in &re_export.items {
161                    // Skip if already added (prefer directly-defined items)
162                    if items.iter().any(|item| matches_item_name(item, item_name)) {
163                        continue;
164                    }
165
166                    // Build re-exported item (link targets will be corrected later)
167                    if let Some(item) = self.build_reexported_item(
168                        &re_export.source_module,
169                        source_module,
170                        item_name,
171                    )? {
172                        items.push(item);
173                    }
174                }
175            }
176        }
177
178        // Process submodules - convert to DocItem::Module entries
179        for submod_name in &module.submodules {
180            if !submod_name.starts_with('_') && exports.contains(submod_name) {
181                // Construct FQN for the submodule
182                let submod_fqn = if name.is_empty() {
183                    submod_name.clone()
184                } else {
185                    format!("{}.{}", name, submod_name)
186                };
187
188                // Retrieve the submodule's doc from stub_info.modules
189                let submod_doc = self
190                    .stub_info
191                    .modules
192                    .get(&submod_fqn)
193                    .map(|m| m.doc.clone())
194                    .unwrap_or_default();
195
196                items.push(DocItem::Module(DocSubmodule {
197                    name: submod_name.clone(),
198                    doc: submod_doc,
199                    fqn: submod_fqn,
200                }));
201            }
202        }
203
204        // Correct all link targets and display text in all items
205        // This ensures both directly-defined and re-exported items have correct references
206        for item in &mut items {
207            self.correct_link_targets(item, name);
208        }
209
210        Ok(DocModule {
211            name: name.to_string(),
212            doc: module.doc.clone(),
213            items,
214            submodules: module
215                .submodules
216                .iter()
217                .filter(|s| !s.starts_with('_'))
218                .cloned()
219                .collect(),
220        })
221    }
222
223    fn build_function(
224        &self,
225        module: &str,
226        func_defs: &[crate::generate::FunctionDef],
227    ) -> Result<DocItem> {
228        // Sort overloads by source location for deterministic ordering (same as stub generation)
229        let mut sorted_defs = func_defs.to_vec();
230        sorted_defs.sort_by_key(|func| (func.file, func.line, func.column, func.index));
231
232        // Requirement #1: Include ALL overload signatures
233        let signatures: Vec<DocSignature> = sorted_defs
234            .iter()
235            .map(|def| self.build_signature(module, def))
236            .collect::<Result<_>>()?;
237
238        // Use first def's doc (they should all have same doc)
239        let doc = sorted_defs
240            .first()
241            .map(|d| d.doc.to_string())
242            .unwrap_or_default();
243
244        let deprecated = sorted_defs.first().and_then(|d| {
245            d.deprecated.as_ref().map(|dep| DeprecatedInfo {
246                since: dep.since.map(|s| s.to_string()),
247                note: dep.note.map(|s| s.to_string()),
248            })
249        });
250
251        Ok(DocItem::Function(DocFunction {
252            name: sorted_defs[0].name.to_string(),
253            doc,
254            signatures,
255            is_async: sorted_defs[0].is_async,
256            deprecated,
257        }))
258    }
259
260    fn build_signature_from_params(
261        &self,
262        module: &str,
263        parameters: &crate::generate::Parameters,
264        return_type: &crate::TypeInfo,
265    ) -> Result<DocSignature> {
266        let ctx = self.create_context(module);
267        let type_renderer = ctx.type_renderer();
268        let default_parser = ctx.default_parser();
269
270        let params: Vec<DocParameter> = parameters
271            .positional_only
272            .iter()
273            .chain(parameters.positional_or_keyword.iter())
274            .chain(parameters.keyword_only.iter())
275            .chain(parameters.varargs.iter())
276            .chain(parameters.varkw.iter())
277            .map(|param| DocParameter {
278                name: param.name.to_string(),
279                type_: type_renderer.render_type(&param.type_info),
280                default: match &param.default {
281                    crate::generate::ParameterDefault::None => None,
282                    crate::generate::ParameterDefault::Expr(s) => {
283                        // Parse default value and identify type references
284                        Some(default_parser.parse(s, &param.type_info))
285                    }
286                },
287            })
288            .collect();
289
290        let ret_type = Some(type_renderer.render_type(return_type));
291
292        Ok(DocSignature {
293            parameters: params,
294            return_type: ret_type,
295        })
296    }
297
298    fn build_signature(
299        &self,
300        module: &str,
301        def: &crate::generate::FunctionDef,
302    ) -> Result<DocSignature> {
303        self.build_signature_from_params(module, &def.parameters, &def.r#return)
304    }
305
306    fn build_signature_from_method(
307        &self,
308        module: &str,
309        def: &crate::generate::MethodDef,
310    ) -> Result<DocSignature> {
311        self.build_signature_from_params(module, &def.parameters, &def.r#return)
312    }
313
314    fn build_type_alias(
315        &self,
316        module: &str,
317        alias: &crate::generate::TypeAliasDef,
318    ) -> Result<DocItem> {
319        // Requirement #2: Preserve alias definition + docstring
320        let ctx = self.create_context(module);
321        let type_renderer = ctx.type_renderer();
322
323        Ok(DocItem::TypeAlias(DocTypeAlias {
324            name: alias.name.to_string(),
325            doc: alias.doc.to_string(),
326            definition: type_renderer.render_type(&alias.type_),
327        }))
328    }
329
330    fn build_class(&self, module: &str, class: &crate::generate::ClassDef) -> Result<DocItem> {
331        let ctx = self.create_context(module);
332        let type_renderer = ctx.type_renderer();
333
334        let bases: Vec<DocTypeExpr> = class
335            .bases
336            .iter()
337            .map(|base| type_renderer.render_type(base))
338            .collect();
339
340        let mut methods = Vec::new();
341        for (method_name, method_overloads) in &class.methods {
342            let signatures: Vec<DocSignature> = method_overloads
343                .iter()
344                .map(|method| self.build_signature_from_method(module, method))
345                .collect::<Result<_>>()?;
346
347            let deprecated = method_overloads.first().and_then(|d| {
348                d.deprecated.as_ref().map(|dep| DeprecatedInfo {
349                    since: dep.since.map(|s| s.to_string()),
350                    note: dep.note.map(|s| s.to_string()),
351                })
352            });
353
354            methods.push(DocFunction {
355                name: method_name.to_string(),
356                doc: method_overloads
357                    .first()
358                    .map(|m| m.doc.to_string())
359                    .unwrap_or_default(),
360                signatures,
361                is_async: method_overloads
362                    .first()
363                    .map(|m| m.is_async)
364                    .unwrap_or(false),
365                deprecated,
366            });
367        }
368
369        let attributes = Vec::new(); // TODO: implement attributes
370
371        Ok(DocItem::Class(DocClass {
372            name: class.name.to_string(),
373            doc: class.doc.to_string(),
374            bases,
375            methods,
376            attributes,
377            deprecated: None, // ClassDef doesn't have deprecated field
378        }))
379    }
380
381    fn build_enum_as_class(
382        &self,
383        _module: &str,
384        enum_def: &crate::generate::EnumDef,
385    ) -> Result<DocItem> {
386        // Convert enum to class-like representation
387        // EnumDef doesn't have bases field
388
389        // Convert enum variants to DocAttributes
390        let attributes: Vec<DocAttribute> = enum_def
391            .variants
392            .iter()
393            .map(|(variant_name, variant_doc)| DocAttribute {
394                name: (*variant_name).to_string(),
395                doc: (*variant_doc).to_string(),
396                type_: None, // Enum variants don't have explicit type annotations
397            })
398            .collect();
399
400        Ok(DocItem::Class(DocClass {
401            name: enum_def.name.to_string(),
402            doc: enum_def.doc.to_string(),
403            bases: Vec::new(), // Enums don't have bases in our structure
404            methods: Vec::new(),
405            attributes,
406            deprecated: None,
407        }))
408    }
409
410    fn build_variable(&self, module: &str, var: &crate::generate::VariableDef) -> Result<DocItem> {
411        let ctx = self.create_context(module);
412        let type_renderer = ctx.type_renderer();
413
414        Ok(DocItem::Variable(DocVariable {
415            name: var.name.to_string(),
416            doc: String::new(), // VariableDef doesn't have doc field
417            type_: Some(type_renderer.render_type(&var.type_)),
418        }))
419    }
420
421    /// Build a re-exported item from a source module
422    fn build_reexported_item(
423        &self,
424        source_module_name: &str,
425        source_module: &crate::generate::Module,
426        item_name: &str,
427    ) -> Result<Option<DocItem>> {
428        // Try functions
429        if let Some(func_defs) = source_module.function.get(item_name) {
430            return Ok(Some(self.build_function(source_module_name, func_defs)?));
431        }
432
433        // Try classes
434        for class_def in source_module.class.values() {
435            if class_def.name == item_name {
436                return Ok(Some(self.build_class(source_module_name, class_def)?));
437            }
438        }
439
440        // Try enums
441        for enum_def in source_module.enum_.values() {
442            if enum_def.name == item_name {
443                return Ok(Some(
444                    self.build_enum_as_class(source_module_name, enum_def)?,
445                ));
446            }
447        }
448
449        // Try type aliases
450        if let Some(alias_def) = source_module.type_aliases.get(item_name) {
451            return Ok(Some(self.build_type_alias(source_module_name, alias_def)?));
452        }
453
454        // Try variables
455        if let Some(var_def) = source_module.variables.get(item_name) {
456            return Ok(Some(self.build_variable(source_module_name, var_def)?));
457        }
458
459        Ok(None)
460    }
461
462    /// Correct link targets in a re-exported item to point to the target module
463    fn correct_link_targets(&self, item: &mut DocItem, _target_module: &str) {
464        match item {
465            DocItem::Function(func) => {
466                for sig in &mut func.signatures {
467                    if let Some(ret) = &mut sig.return_type {
468                        self.correct_type_expr(ret);
469                    }
470                    for param in &mut sig.parameters {
471                        self.correct_type_expr(&mut param.type_);
472                    }
473                }
474            }
475            DocItem::Class(cls) => {
476                for base in &mut cls.bases {
477                    self.correct_type_expr(base);
478                }
479                for method in &mut cls.methods {
480                    for sig in &mut method.signatures {
481                        if let Some(ret) = &mut sig.return_type {
482                            self.correct_type_expr(ret);
483                        }
484                        for param in &mut sig.parameters {
485                            self.correct_type_expr(&mut param.type_);
486                        }
487                    }
488                }
489                for attr in &mut cls.attributes {
490                    if let Some(type_) = &mut attr.type_ {
491                        self.correct_type_expr(type_);
492                    }
493                }
494            }
495            DocItem::TypeAlias(alias) => {
496                self.correct_type_expr(&mut alias.definition);
497            }
498            DocItem::Variable(var) => {
499                if let Some(type_) = &mut var.type_ {
500                    self.correct_type_expr(type_);
501                }
502            }
503            DocItem::Module(_) => {}
504        }
505    }
506
507    /// Correct a type expression to use export_map for link targets
508    fn correct_type_expr(&self, type_expr: &mut DocTypeExpr) {
509        if let Some(link_target) = &mut type_expr.link_target {
510            // Try to find the correct export module and FQN
511            // First try the original FQN
512            let (exported_fqn, exported_module) =
513                if let Some(module) = self.export_map.get(&link_target.fqn) {
514                    (link_target.fqn.clone(), module.clone())
515                } else {
516                    // If not found, try to extract the type name and look for it under other modules
517                    // e.g., "hidden_module_docgen_test._core.A" -> try "hidden_module_docgen_test.A"
518                    if let Some(type_name) = link_target.fqn.split('.').next_back() {
519                        // Try each module in export_map to find a match
520                        if let Some((fqn, module)) = self
521                            .export_map
522                            .iter()
523                            .find(|(fqn, _)| fqn.ends_with(&format!(".{}", type_name)))
524                        {
525                            (fqn.clone(), module.clone())
526                        } else {
527                            (link_target.fqn.clone(), link_target.doc_module.clone())
528                        }
529                    } else {
530                        (link_target.fqn.clone(), link_target.doc_module.clone())
531                    }
532                };
533
534            link_target.fqn = exported_fqn;
535            link_target.doc_module = exported_module;
536
537            // Update display text to strip internal module prefixes
538            // e.g., "_core.A" -> "A"
539            type_expr.display = self.strip_internal_module_prefix(&type_expr.display);
540        }
541        for child in &mut type_expr.children {
542            self.correct_type_expr(child);
543        }
544    }
545
546    /// Strip internal module prefixes (modules starting with '_') from display text
547    /// e.g., "_core.A" -> "A", "_internal.Foo" -> "Foo"
548    fn strip_internal_module_prefix(&self, display: &str) -> String {
549        prefix_stripper::strip_internal_prefixes(display)
550    }
551}
pyo3_stub_gen/docgen/builder.rs

pyo3_stub_gen/docgen/
builder.rs
