pyo3_stub_gen/generate/

module.rs

use crate::generate::*;
use crate::stub_type::ImportRef;
use itertools::Itertools;
use std::{
    any::TypeId,
    collections::{BTreeMap, BTreeSet},
    fmt,
};

/// Re-export from another module for __all__
#[derive(Debug, Clone, PartialEq)]
pub struct ModuleReExport {
    pub source_module: String,
    /// Items to re-export. Empty means wildcard (will be resolved).
    pub items: Vec<String>,
    /// Additional items to include with wildcard (e.g., `__version__`).
    /// These are merged into `items` after wildcard resolution.
    pub additional_items: Vec<String>,
}

/// Type info for a Python (sub-)module. This corresponds to a single `*.pyi` file.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct Module {
    pub doc: String,
    pub class: BTreeMap<TypeId, ClassDef>,
    pub enum_: BTreeMap<TypeId, EnumDef>,
    pub function: BTreeMap<&'static str, Vec<FunctionDef>>,
    pub variables: BTreeMap<&'static str, VariableDef>,
    pub type_aliases: BTreeMap<&'static str, TypeAliasDef>,
    pub name: String,
    pub default_module_name: String,
    /// Direct submodules of this module.
    pub submodules: BTreeSet<String>,
    /// Module re-exports for __all__
    pub module_re_exports: Vec<ModuleReExport>,
    /// Verbatim entries to add to __all__
    pub verbatim_all_entries: BTreeSet<String>,
    /// Explicitly excluded entries from __all__
    pub excluded_all_entries: BTreeSet<String>,
}

impl Module {
    /// Check if this module has no content to generate.
    ///
    /// Returns true if the module has no classes, enums, functions, variables,
    /// type aliases, submodules, re-exports, docstrings, or verbatim entries.
    /// Modules that are empty should be skipped during generation.
    pub fn is_empty(&self) -> bool {
        self.doc.is_empty()
            && self.class.is_empty()
            && self.enum_.is_empty()
            && self.function.is_empty()
            && self.variables.is_empty()
            && self.type_aliases.is_empty()
            && self.submodules.is_empty()
            && self.module_re_exports.is_empty()
            && self.verbatim_all_entries.is_empty()
    }

    /// Check if this module can have `__init__.py` generated.
    ///
    /// Returns true if the module has no PyO3-generated items (classes, enums,
    /// functions, variables, type aliases). Such modules can only contain
    /// re-exports and docstrings, which can be represented in `__init__.py`.
    pub fn is_init_py_compatible(&self) -> bool {
        self.class.is_empty()
            && self.enum_.is_empty()
            && self.function.is_empty()
            && self.variables.is_empty()
            && self.type_aliases.is_empty()
    }

    /// Get the names of all declared items in this module.
    ///
    /// Returns a list of item names that were declared via `gen_stub_*` macros.
    pub fn declared_item_names(&self) -> Vec<String> {
        let mut names = Vec::new();

        if !self.doc.is_empty() {
            names.push("module_doc".to_string());
        }
        for class in self.class.values() {
            names.push(format!("class {}", class.name));
        }
        for enum_def in self.enum_.values() {
            names.push(format!("enum {}", enum_def.name));
        }
        for func_name in self.function.keys() {
            names.push(format!("function {}", func_name));
        }
        for var_name in self.variables.keys() {
            names.push(format!("variable {}", var_name));
        }
        for alias_name in self.type_aliases.keys() {
            names.push(format!("type_alias {}", alias_name));
        }
        for re_export in &self.module_re_exports {
            names.push(format!("re-export from {}", re_export.source_module));
        }

        names.sort();
        names
    }

    /// Format module with configuration for type alias syntax, returning a String
    pub fn format_with_config(&self, use_type_statement: bool) -> String {
        use std::fmt::Write;
        let mut output = String::new();

        // Use a custom formatter struct
        struct ModuleFormatter<'a> {
            module: &'a Module,
            use_type_statement: bool,
        }

        impl<'a> fmt::Display for ModuleFormatter<'a> {
            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                // Write header and docstring
                writeln!(f, "# This file is automatically generated by pyo3_stub_gen")?;
                writeln!(f, "# ruff: noqa: E501, F401, F403, F405")?;
                if !self.module.doc.is_empty() {
                    docstring::write_docstring(f, &self.module.doc, "")?;
                }
                writeln!(f)?;

                let mut imports = self.module.import();

                // Conditionally add TypeAlias import
                if !self.use_type_statement && !self.module.type_aliases.is_empty() {
                    imports.insert(ImportRef::Type(crate::stub_type::TypeRef {
                        module: crate::stub_type::ModuleRef::Named("typing".to_string()),
                        name: "TypeAlias".to_string(),
                    }));
                }

                // Check for overload decorator
                let any_overloaded = self.module.function.values().any(|functions| {
                    let has_overload = functions.iter().any(|func| func.is_overload);
                    functions.len() > 1 && has_overload
                });
                if any_overloaded {
                    imports.insert("typing".into());
                }

                // Generate imports (same logic as Display impl)
                let mut type_ref_grouped: BTreeMap<String, Vec<String>> = BTreeMap::new();
                for import_ref in imports.into_iter().sorted() {
                    match import_ref {
                        ImportRef::Module(module_ref) => {
                            let name = module_ref.get().unwrap_or(&self.module.default_module_name);
                            if name != self.module.name && !name.is_empty() {
                                let is_internal_module = if let Some(root) =
                                    self.module.default_module_name.split('.').next()
                                {
                                    name.starts_with(root)
                                } else {
                                    false
                                };

                                if is_internal_module && name.contains('.') {
                                    let last_dot_pos = name.rfind('.').unwrap();
                                    let parent_module = &name[..last_dot_pos];
                                    let child_module = &name[last_dot_pos + 1..];

                                    if !self.module.submodules.contains(child_module) {
                                        writeln!(
                                            f,
                                            "from {} import {}",
                                            parent_module, child_module
                                        )?;
                                    }
                                } else {
                                    writeln!(f, "import {name}")?;
                                }
                            }
                        }
                        ImportRef::Type(type_ref) => {
                            let module_name = type_ref
                                .module
                                .get()
                                .unwrap_or(&self.module.default_module_name);
                            if module_name != self.module.name {
                                type_ref_grouped
                                    .entry(module_name.to_string())
                                    .or_default()
                                    .push(type_ref.name);
                            }
                        }
                    }
                }
                for (module_name, type_names) in type_ref_grouped {
                    let mut sorted_type_names = type_names.clone();
                    sorted_type_names.sort();
                    writeln!(
                        f,
                        "from {} import {}",
                        module_name,
                        sorted_type_names.join(", ")
                    )?;
                }

                // Add imports for module re-exports (always explicit, not wildcard)
                let mut sorted_re_exports = self.module.module_re_exports.clone();
                sorted_re_exports.sort_by(|a, b| a.source_module.cmp(&b.source_module));
                for re_export in &sorted_re_exports {
                    if re_export.items.is_empty() {
                        continue;
                    }
                    let mut sorted_items = re_export.items.clone();
                    sorted_items.sort();
                    writeln!(
                        f,
                        "from {} import {}",
                        re_export.source_module,
                        sorted_items.join(", ")
                    )?;
                }
                for submod in &self.module.submodules {
                    writeln!(f, "from . import {submod}")?;
                }

                // Generate __all__ list
                self.module.write_all_list(f)?;

                writeln!(f)?;

                // Generate type aliases with configuration
                for alias in self.module.type_aliases.values() {
                    alias.fmt_with_config(&self.module.name, f, self.use_type_statement)?;
                    writeln!(f)?;
                }

                // Generate variables
                for var in self.module.variables.values() {
                    var.fmt_for_module(&self.module.name, f)?;
                    writeln!(f)?;
                }

                // Generate classes
                for class in self.module.class.values().sorted_by_key(|class| class.name) {
                    class.fmt_for_module(&self.module.name, f)?;
                }

                // Generate enums
                for enum_ in self.module.enum_.values().sorted_by_key(|enum_| enum_.name) {
                    enum_.fmt_for_module(&self.module.name, f)?;
                }

                // Generate functions
                for functions in self.module.function.values() {
                    let has_overload = functions.iter().any(|func| func.is_overload);
                    let should_add_overload = functions.len() > 1 && has_overload;

                    let mut sorted_functions = functions.clone();
                    sorted_functions
                        .sort_by_key(|func| (func.file, func.line, func.column, func.index));
                    for function in sorted_functions {
                        if should_add_overload {
                            writeln!(f, "@typing.overload")?;
                        }
                        function.fmt_for_module(&self.module.name, f)?;
                    }
                }

                Ok(())
            }
        }

        write!(
            &mut output,
            "{}",
            ModuleFormatter {
                module: self,
                use_type_statement
            }
        )
        .unwrap();
        output
    }

    /// Collect all items for the `__all__` list in `.pyi` stub files.
    ///
    /// This collects public items from classes, enums, functions, variables,
    /// type aliases, submodules, re-exports, and verbatim entries.
    /// Items starting with `_` are excluded unless added via `export_verbatim!`.
    fn collect_all_items(&self) -> BTreeSet<String> {
        let mut all_items: BTreeSet<String> = BTreeSet::new();

        // Collect public items from this module
        for class in self.class.values() {
            if !class.name.starts_with('_') {
                all_items.insert(class.name.to_string());
            }
        }
        for enum_ in self.enum_.values() {
            if !enum_.name.starts_with('_') {
                all_items.insert(enum_.name.to_string());
            }
        }
        for func_name in self.function.keys() {
            if !func_name.starts_with('_') {
                all_items.insert(func_name.to_string());
            }
        }
        for var_name in self.variables.keys() {
            if !var_name.starts_with('_') {
                all_items.insert(var_name.to_string());
            }
        }
        for alias_name in self.type_aliases.keys() {
            if !alias_name.starts_with('_') {
                all_items.insert(alias_name.to_string());
            }
        }
        for submod in &self.submodules {
            if !submod.starts_with('_') {
                all_items.insert(submod.to_string());
            }
        }

        // Add items from re-exports
        for re_export in &self.module_re_exports {
            all_items.extend(re_export.items.iter().cloned());
        }

        // Add verbatim entries (allows explicit inclusion of underscore items)
        all_items.extend(self.verbatim_all_entries.iter().cloned());

        // Remove explicitly excluded items
        for excluded in &self.excluded_all_entries {
            all_items.remove(excluded);
        }

        all_items
    }

    /// Format module as `__init__.py` content.
    ///
    /// This generates a Python `__init__.py` file with a module docstring,
    /// re-exports, and an `__all__` list.
    /// Unlike `format_with_config()` which generates `.pyi` stub files, this generates
    /// actual Python code for runtime use.
    pub fn format_init_py(&self) -> String {
        use std::fmt::Write;
        let mut output = String::new();

        // Header
        writeln!(
            output,
            "# This file is automatically generated by pyo3_stub_gen"
        )
        .unwrap();
        writeln!(output, "# ruff: noqa: F401").unwrap();
        docstring::write_docstring(&mut output, &self.doc, "").unwrap();
        writeln!(output).unwrap();

        // Re-export imports (sorted for deterministic output)
        // Always use explicit imports (not wildcard) for better tooling support
        let mut sorted_re_exports = self.module_re_exports.clone();
        sorted_re_exports.sort_by(|a, b| a.source_module.cmp(&b.source_module));
        for re_export in &sorted_re_exports {
            if re_export.items.is_empty() {
                continue;
            }
            let mut sorted_items = re_export.items.clone();
            sorted_items.sort();
            writeln!(
                output,
                "from {} import {}",
                re_export.source_module,
                sorted_items.join(", ")
            )
            .unwrap();
        }

        // Collect __all__ items from re-exports only
        // (Unlike __init__.pyi, __init__.py can only contain re-exported items)
        let mut all_items: BTreeSet<String> = BTreeSet::new();
        for re_export in &self.module_re_exports {
            all_items.extend(re_export.items.iter().cloned());
        }
        all_items.extend(self.verbatim_all_entries.iter().cloned());
        for excluded in &self.excluded_all_entries {
            all_items.remove(excluded);
        }
        if all_items.is_empty() {
            writeln!(output, "__all__ = []").unwrap();
        } else {
            writeln!(output, "__all__ = [").unwrap();
            for item in all_items {
                writeln!(output, "    \"{}\",", item).unwrap();
            }
            writeln!(output, "]").unwrap();
        }

        output
    }

    fn write_all_list(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let all_items = self.collect_all_items();

        // Always write __all__ list (even if empty for consistency)
        if all_items.is_empty() {
            writeln!(f, "__all__ = []")?;
        } else {
            writeln!(f, "__all__ = [")?;
            for item in all_items {
                writeln!(f, "    \"{}\",", item)?;
            }
            writeln!(f, "]")?;
        }

        Ok(())
    }
}

impl Import for Module {
    fn import(&self) -> HashSet<ImportRef> {
        let mut imports = HashSet::new();
        for class in self.class.values() {
            imports.extend(class.import());
        }
        for enum_ in self.enum_.values() {
            imports.extend(enum_.import());
        }
        for function in self.function.values().flatten() {
            imports.extend(function.import());
        }
        for variable in self.variables.values() {
            imports.extend(variable.import());
        }
        for type_alias in self.type_aliases.values() {
            imports.extend(type_alias.import());
        }
        imports
    }
}

impl fmt::Display for Module {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // Write header and docstring directly (no type names here)
        writeln!(f, "# This file is automatically generated by pyo3_stub_gen")?;
        writeln!(f, "# ruff: noqa: E501, F401, F403, F405")?;
        if !self.doc.is_empty() {
            docstring::write_docstring(f, &self.doc, "")?;
        }
        writeln!(f)?;

        let mut imports = self.import();
        // Check if any function group needs @overload decorator
        let any_overloaded = self.function.values().any(|functions| {
            let has_overload = functions.iter().any(|f| f.is_overload);
            functions.len() > 1 && has_overload
        });
        if any_overloaded {
            imports.insert("typing".into());
        }

        // To gather `from submod import A, B, C` style imports
        let mut type_ref_grouped: BTreeMap<String, Vec<String>> = BTreeMap::new();
        for import_ref in imports.into_iter().sorted() {
            match import_ref {
                ImportRef::Module(module_ref) => {
                    let name = module_ref.get().unwrap_or(&self.default_module_name);
                    if name != self.name && !name.is_empty() {
                        // Check if this is a module within the current package
                        // by checking if the module name starts with the package name
                        let is_internal_module =
                            if let Some(root) = self.default_module_name.split('.').next() {
                                name.starts_with(root)
                            } else {
                                false
                            };

                        // For nested modules like "package.module.submodule" within the current package
                        // Generate: from package.module import submodule
                        // For external modules like "collections.abc", use: import collections.abc
                        if is_internal_module && name.contains('.') {
                            let last_dot_pos = name.rfind('.').unwrap();
                            let parent_module = &name[..last_dot_pos];
                            let child_module = &name[last_dot_pos + 1..];

                            // Skip if this is a direct submodule (already imported via submodule imports)
                            if !self.submodules.contains(child_module) {
                                writeln!(f, "from {} import {}", parent_module, child_module)?;
                            }
                        } else {
                            // External module or top-level module - use standard import
                            writeln!(f, "import {name}")?;
                        }
                    }
                }
                ImportRef::Type(type_ref) => {
                    let module_name = type_ref.module.get().unwrap_or(&self.default_module_name);
                    if module_name != self.name {
                        type_ref_grouped
                            .entry(module_name.to_string())
                            .or_default()
                            .push(type_ref.name);
                    }
                }
            }
        }
        for (module_name, type_names) in type_ref_grouped {
            let mut sorted_type_names = type_names.clone();
            sorted_type_names.sort();
            writeln!(
                f,
                "from {} import {}",
                module_name,
                sorted_type_names.join(", ")
            )?;
        }
        // Add imports for module re-exports (always explicit, not wildcard)
        let mut sorted_re_exports = self.module_re_exports.clone();
        sorted_re_exports.sort_by(|a, b| a.source_module.cmp(&b.source_module));
        for re_export in &sorted_re_exports {
            if re_export.items.is_empty() {
                continue;
            }
            let mut sorted_items = re_export.items.clone();
            sorted_items.sort();
            writeln!(
                f,
                "from {} import {}",
                re_export.source_module,
                sorted_items.join(", ")
            )?;
        }
        for submod in &self.submodules {
            writeln!(f, "from . import {submod}")?;
        }

        // Generate __all__ list
        self.write_all_list(f)?;

        writeln!(f)?;

        for alias in self.type_aliases.values() {
            alias.fmt_for_module(&self.name, f)?;
            writeln!(f)?;
        }
        for var in self.variables.values() {
            var.fmt_for_module(&self.name, f)?;
            writeln!(f)?;
        }
        for class in self.class.values().sorted_by_key(|class| class.name) {
            class.fmt_for_module(&self.name, f)?;
        }
        for enum_ in self.enum_.values().sorted_by_key(|class| class.name) {
            enum_.fmt_for_module(&self.name, f)?;
        }
        for functions in self.function.values() {
            // Check if we should add @overload to all functions
            let has_overload = functions.iter().any(|func| func.is_overload);
            let should_add_overload = functions.len() > 1 && has_overload;

            // Sort by source location and index for deterministic ordering
            let mut sorted_functions = functions.clone();
            sorted_functions.sort_by_key(|func| (func.file, func.line, func.column, func.index));
            for function in sorted_functions {
                if should_add_overload {
                    writeln!(f, "@typing.overload")?;
                }
                function.fmt_for_module(&self.name, f)?;
            }
        }
        Ok(())
    }
}