generate_mkdocs.py

"""
@file: generate_mkdocs.py
@breif: Used to automatically generate MkDocs documentation for Python libraries.
@author: Wu Maojia
@update: 2025.10.3
"""
import os
import ast
import yaml
import shutil


def extract_classes(file_path: str):
    """
    Extract the names of all classes in a Python file.

    Parameters:
        file_path (str): Path to the Python file.

    Returns:
        class_names (list): List of class names.
    """
    class_names = []
    with open(file_path, 'r', encoding='utf-8') as f:
        node = ast.parse(f.read(), filename=file_path)
    for child in node.body:
        if isinstance(child, ast.ClassDef):
            class_names.append(child.name)
    return class_names


def generate_api_docs(root_folder: str, output_folder: str, index_file: str, mkdocs_file: str, ex_home_file: str, ex_assets_folder: str):
    """
    Automatically generate Markdown files for API documentation, update the homepage, and modify mkdocs.yml.

    Parameters:
        root_folder (str): Path to the root directory of the library.
        output_folder (str): Path to the output directory for the generated documentation.
        index_file (str): Path to the homepage file.
        mkdocs_file (str): Path to the mkdocs.yml file.
    """
    nav_structure = {}
    tutorials_structure = {}

    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # Load existing mkdocs config to preserve Tutorials section if it exists
    existing_tutorials_nav = None
    if os.path.exists(mkdocs_file):
        with open(mkdocs_file, 'r', encoding='utf-8') as f:
            existing_config = yaml.unsafe_load(f)
            if existing_config and 'nav' in existing_config:
                for item in existing_config['nav']:
                    if isinstance(item, dict) and 'Tutorials' in item:
                        existing_tutorials_nav = item['Tutorials']
                        break

    # Process Python source code for API documentation
    for root, dirs, files in os.walk(root_folder):
        relative_root = os.path.relpath(root, root_folder)  # Get the relative path of the current directory
        for file in files:
            if file.endswith('.py') and not file.startswith('__'):
                file_path = os.path.join(root, file)
                class_names = extract_classes(file_path)
                if not class_names:
                    print(f"Warning: No class found in {file_path}. Skipping...")
                    continue

                module_path = file_path.replace('.py', '').replace('/', '.').replace('\\', '.')
                # Create a corresponding output directory that mirrors the input directory's structure
                output_dir = os.path.join(output_folder, 'api', relative_root, file.replace('.py', ''))
                os.makedirs(output_dir, exist_ok=True)

                # Generate Markdown files for each class
                for class_name in class_names:
                    class_md_path = os.path.join(output_dir, f"{class_name}.md")
                    with open(class_md_path, 'w', encoding='utf-8') as f:
                        f.write(f"# {class_name}\n\n::: {module_path}.{class_name}\n")

                # Update navigation structure
                current_nav = nav_structure
                for part in relative_root.split(os.sep):
                    if part:
                        current_nav = current_nav.setdefault(part, {})
                current_nav.setdefault(file.replace('.py', ''), []).extend([{
                    class_name: os.path.relpath(os.path.join(output_dir, f"{class_name}.md"), output_folder)
                } for class_name in class_names])

    # Process tutorials folder only if no existing Tutorials nav found
    tutorials_folder = 'tutorials'
    if existing_tutorials_nav is None and os.path.exists(tutorials_folder):
        # Copy tutorials folder to docs
        tutorials_output = os.path.join(output_folder, 'tutorials')
        if os.path.exists(tutorials_output):
            shutil.rmtree(tutorials_output)
        shutil.copytree(tutorials_folder, tutorials_output)
        
        # Build tutorials navigation structure
        for root, dirs, files in os.walk(tutorials_folder):
            relative_root = os.path.relpath(root, tutorials_folder)
            for file in files:
                if file.endswith('.md'):
                    file_path = os.path.join(root, file)
                    relative_path = os.path.relpath(file_path, tutorials_folder)
                    output_path = os.path.join('tutorials', relative_path).replace('\\', '/')
                    
                    # Build nested structure
                    current_nav = tutorials_structure
                    parts = relative_path.split(os.sep)
                    
                    # Handle nested directories
                    for part in parts[:-1]:  # All parts except the filename
                        if part:
                            current_nav = current_nav.setdefault(part, {})
                    
                    # Add the markdown file
                    file_name_without_ext = os.path.splitext(parts[-1])[0]
                    if isinstance(current_nav, dict):
                        current_nav[file_name_without_ext] = output_path
                    else:
                        current_nav = {file_name_without_ext: output_path}

    # Copy assets folder if exists
    if os.path.exists(ex_assets_folder):
        assets_output = os.path.join(output_folder, ex_assets_folder)
        if os.path.exists(assets_output):
            shutil.rmtree(assets_output)
        shutil.copytree(ex_assets_folder, assets_output)

    # Update homepage
    with open(ex_home_file, 'r', encoding='utf-8') as f:
        ex_home_content = f.read()

    with open(index_file, 'w', encoding='utf-8') as f:
        f.write(ex_home_content)

    # Build the nav section of mkdocs.yml
    nav = [{"Home": "index.md"}]
    
    # Add tutorials to navigation if they exist (preserve existing if available)
    if existing_tutorials_nav is not None:
        nav.append({"Tutorials": existing_tutorials_nav})
    elif tutorials_structure:
        nav.append({"Tutorials": build_nav_structure(tutorials_structure)})
    
    # Add API reference to navigation
    nav.append({"API Reference": build_nav_structure(nav_structure)})

    print("\nGenerated nav for mkdocs.yml:")
    print(yaml.dump({"nav": nav}, allow_unicode=True, sort_keys=False))

    # If mkdocs.yml exists, automatically update its nav section
    if os.path.exists(mkdocs_file):
        with open(mkdocs_file, 'r', encoding='utf-8') as f:
            mkdocs_config = yaml.unsafe_load(f)

        mkdocs_config['nav'] = nav

        with open(mkdocs_file, 'w', encoding='utf-8') as f:
            yaml.dump(mkdocs_config, f, allow_unicode=True, sort_keys=False)

def build_nav_structure(nav_dict: dict) -> list:
    """
    Build navigation structure from dictionary.
    
    Parameters:
        nav_dict (dict): Navigation dictionary
        
    Returns:
        list: Navigation structure for mkdocs
    """
    result = []
    for key, value in sorted(nav_dict.items()):
        if isinstance(value, dict):
            # It's a category with subitems
            result.append({key: build_nav_structure(value)})
        elif isinstance(value, list):
            # It's a list of class items (for API reference)
            class_items = []
            for item in value:
                if isinstance(item, dict):
                    for class_name, doc_path in item.items():
                        class_items.append({class_name: doc_path})
            result.append({key: class_items})
        else:
            # It's a direct file reference (for tutorials)
            result.append({key: value})
    return result


if __name__ == '__main__':
    # Example usage
    generate_api_docs(
        root_folder='src/python_motion_planning',  # Code directory
        output_folder='docs/',  # Directory for the generated documentation
        index_file='docs/index.md',  # Path to the homepage file
        mkdocs_file='mkdocs.yml',  # Path to the mkdocs.yml file
        ex_home_file='README.md',   # Extern homepage file
        ex_assets_folder='assets/'  # Extern assets folder
    )