build.py

#!/usr/bin/env python3
"""
Build script to combine multiple Python source files from src/ into build/visualization.py
Uses AST for dependency analysis and ruff for formatting.
"""

import ast
import subprocess
from pathlib import Path

import networkx as nx
import toml


class ImportVisitor(ast.NodeVisitor):
    """AST visitor to extract imports from Python files."""

    def __init__(self):
        self.standard_imports = set()
        self.from_imports = {}  # module -> set of imported names

    def visit_Import(self, node):
        for alias in node.names:
            # Skip relative imports
            if not alias.name.startswith("."):
                self.standard_imports.add(alias.name)
        self.generic_visit(node)

    def visit_ImportFrom(self, node):
        # Process all non-relative imports (from module import name)
        if node.level == 0 and node.module:
            # This is an external import that we need to preserve
            imported_names = {alias.name for alias in node.names}
            if imported_names:
                if node.module not in self.from_imports:
                    self.from_imports[node.module] = set()
                self.from_imports[node.module].update(imported_names)
        # Track sibling imports for dependency analysis
        elif node.level == 1 and node.module:
            module_name = node.module
            imported_names = {alias.name for alias in node.names}
            self.from_imports[module_name] = imported_names
        self.generic_visit(node)


def analyze_file_imports(file_path: Path) -> tuple[set[str], dict[str, set[str]], dict[str, set[str]]]:
    """Analyze imports from a Python file using AST. Returns (standard_imports, from_imports, sibling_imports)."""
    try:
        with open(file_path, encoding="utf-8") as f:
            tree = ast.parse(f.read())

        visitor = ImportVisitor()
        visitor.visit(tree)

        # Separate sibling imports from external imports
        sibling_imports = {}
        external_imports = {}

        for module, names in visitor.from_imports.items():
            # Check if this is a sibling import by seeing if any file in the current directory matches
            if (file_path.parent / f"{module}.py").exists():
                sibling_imports[module] = names
            else:
                external_imports[module] = names

        return visitor.standard_imports, external_imports, sibling_imports
    except Exception as e:
        print(f"Warning: Could not parse {file_path}: {e}")
        return set(), {}, {}


def build_dependency_graph(src_files: list[Path]) -> nx.DiGraph:
    """Build dependency graph using networkx."""
    G = nx.DiGraph()

    # Add nodes
    for f in src_files:
        G.add_node(f.name, path=f)

    # Add edges based on sibling imports
    for f in src_files:
        _, _, sibling_imports = analyze_file_imports(f)
        for imported_module in sibling_imports.keys():
            # Find the corresponding file
            for other_file in src_files:
                if other_file.stem == imported_module:
                    G.add_edge(f.name, other_file.name)
                    break

    return G


def get_build_order(src_files: list[Path]) -> list[Path]:
    """Get files in dependency order using topological sort."""
    if not src_files:
        return []

    G = build_dependency_graph(src_files)

    # Check for circular dependencies
    if not nx.is_directed_acyclic_graph(G):
        print("⚠️  Warning: Circular dependencies detected, using alphabetical order as fallback")
        return sorted(src_files, key=lambda p: p.name)

    try:
        ordered_names = list(nx.topological_sort(G))
        # Reverse the order so dependencies come before dependents
        ordered_names.reverse()
        return [G.nodes[n]["path"] for n in ordered_names]
    except nx.NetworkXError:
        return sorted(src_files, key=lambda p: p.name)


def extract_external_imports(src_files: list[Path]) -> list[str]:
    """Extract all external imports (non-relative) from source files."""
    all_imports = set()
    all_from_imports = {}

    for src_file in src_files:
        imports, from_imports, _ = analyze_file_imports(src_file)
        all_imports.update(imports)

        # Merge from imports
        for module, names in from_imports.items():
            if module not in all_from_imports:
                all_from_imports[module] = set()
            all_from_imports[module].update(names)

    # Convert to import statements - let ruff handle the sorting
    import_statements = [f"import {imp}" for imp in all_imports if not imp.startswith(".")]

    # Add from-import statements
    for module in sorted(all_from_imports.keys()):
        names = sorted(all_from_imports[module])
        import_statements.append(f"from {module} import {', '.join(names)}")

    return import_statements


def get_dependencies_from_pyproject() -> list[str]:
    """Extract dependencies from pyproject.toml, excluding open-webui."""
    pyproject_path = Path(__file__).parent / "pyproject.toml"
    with open(pyproject_path) as f:
        pyproject_data = toml.load(f)

    dependencies = pyproject_data.get("project", {}).get("dependencies", [])

    # Filter out open-webui and dev-only dependencies
    filtered_deps = []
    for dep in dependencies:
        dep_name = dep.split(">=")[0].split("==")[0].split(">")[0].split("<")[0].strip()
        if dep_name not in ["open-webui", "ruff", "networkx", "toml"]:
            filtered_deps.append(dep)

    print(f"📦 Extracted {len(filtered_deps)} dependencies from pyproject.toml")
    return filtered_deps


def get_project_metadata() -> dict[str, str]:
    """Extract project metadata from pyproject.toml."""
    pyproject_path = Path(__file__).parent / "pyproject.toml"
    with open(pyproject_path) as f:
        pyproject_data = toml.load(f)
    
    project_data = pyproject_data.get("project", {})
    
    # Extract author name from authors array if available
    author = "Unknown"
    authors = project_data.get("authors", [])
    if authors and isinstance(authors, list) and len(authors) > 0:
        author = authors[0].get("name", "Unknown")
    elif "author" in project_data:
        author = project_data["author"]
    
    return {
        "name": project_data.get("name", "Visualizations"),
        "version": project_data.get("version", "1.0.0"),
        "description": project_data.get("description", "Visualizations for OpenWebUI"),
        "author": author
    }


def concatenate_files(src_files: list[Path], output_file: Path):
    """Concatenate source files with heavy comment headers."""
    output_parts = []

    # Get metadata from pyproject.toml
    metadata = get_project_metadata()
    
    # Add file header with dynamic metadata
    output_parts.append('"""')
    output_parts.append(f"title: {metadata['name']}")
    output_parts.append(f"description: {metadata['description']}")
    output_parts.append(f"author: {metadata['author']}")
    output_parts.append(f"version: {metadata['version']}")

    # Dynamic requirements from pyproject.toml
    dependencies = get_dependencies_from_pyproject()
    requirements_str = ", ".join(dependencies)
    output_parts.append(f"requirements: {requirements_str}")
    output_parts.append('"""')
    output_parts.append("")

    # Extract and add external imports - let ruff handle ordering
    external_imports = extract_external_imports(src_files)
    if external_imports:
        output_parts.extend(external_imports)
        output_parts.append("")

    # Process each source file and add content (excluding imports)
    for src_file in src_files:
        print(f"  📄 Processing {src_file.name}")

        with open(src_file, encoding="utf-8") as f:
            content = f.read()

        # Parse and filter out imports
        try:
            tree = ast.parse(content)

            # Get line numbers of import statements to skip them
            import_lines = set()
            for node in ast.walk(tree):
                if isinstance(node, (ast.Import, ast.ImportFrom)):
                    import_lines.update(range(node.lineno, node.end_lineno + 1))

            # Filter content lines
            lines = content.split("\n")
            filtered_lines = []

            for i, line in enumerate(lines, 1):
                # Skip import lines and sibling import lines
                if i in import_lines:
                    continue
                if line.strip().startswith("from ."):
                    # Add comment about sibling import
                    module_name = line.strip().split(".")[1].split()[0]
                    filtered_lines.append(f"# Functions from {module_name} are available above in this combined file")
                    continue
                filtered_lines.append(line)

            # Remove any leading empty lines
            while filtered_lines and not filtered_lines[0].strip():
                filtered_lines.pop(0)

            # Add section header if content exists
            if filtered_lines:
                file_header = f"\n\n# ================================================================\n# {src_file.name}\n# ================================================================\n"
                output_parts.append(file_header + "\n".join(filtered_lines))

        except Exception as e:
            print(f"  ⚠️  Warning: Could not parse {src_file.name}: {e}")
            # Fallback: add raw content with simple processing
            lines = content.split("\n")
            non_import_lines = []
            skip_imports = True

            for line in lines:
                stripped = line.strip()
                # Stop skipping when we hit non-import, non-comment code
                if skip_imports and stripped and not stripped.startswith("#") and not stripped.startswith(("import ", "from ")):
                    skip_imports = False

                # Skip imports and sibling imports
                if skip_imports and stripped.startswith(("import ", "from ")):
                    continue
                if stripped.startswith("from ."):
                    continue

                non_import_lines.append(line)

            if non_import_lines:
                file_header = f"\n\n# ================================================================\n# {src_file.name}\n# ================================================================\n"
                output_parts.append(file_header + "\n".join(non_import_lines))

    # Write concatenated file
    output_file.parent.mkdir(parents=True, exist_ok=True)
    with open(output_file, "w", encoding="utf-8") as f:
        f.write("\n".join(output_parts))


def run_ruff_format(output_file: Path):
    """Format the output file using ruff."""
    try:
        print("  🔧 Formatting with ruff...")
        # Format the code
        result = subprocess.run(["ruff", "format", str(output_file)], capture_output=True, text=True)
        if result.returncode != 0:
            print(f"  ⚠️  Ruff format warnings: {result.stderr}")

        # Also run ruff check --fix to fix lint issues
        result = subprocess.run(["ruff", "check", "--fix", str(output_file)], capture_output=True, text=True)
        if result.returncode != 0:
            print(f"  ⚠️  Ruff check warnings: {result.stderr}")

        print("  ✅ Ruff formatting completed")
    except FileNotFoundError:
        print("  ⚠️  Ruff not found. Please install with: pip install ruff")
        raise


def check_ruff_available():
    """Check if ruff is available."""
    try:
        subprocess.run(["ruff", "--version"], capture_output=True, check=True)
        return True
    except (subprocess.CalledProcessError, FileNotFoundError):
        return False


def main():
    """Main build function"""
    root = Path(__file__).parent
    src_dir = root / "src"
    build_dir = root / "build"
    output_file = build_dir / "visualization.py"

    print("=" * 70)
    print("🏗️  Building visualization.py")
    print("=" * 70)

    # Check for ruff
    if not check_ruff_available():
        print("⚠️  Ruff not found. Please install it with: pip install ruff")
        print("   Or install dev dependencies with: uv sync")
        return

    # Verify source directory
    if not src_dir.exists():
        print(f"\n❌ Source directory not found: {src_dir}")
        return

    # Get source python files
    src_files = list(src_dir.glob("*.py"))
    if not src_files:
        print(f"\n❌ No Python files found in {src_dir}")
        return

    # Get files in dependency order
    ordered_files = get_build_order(src_files)
    print(f"\n📦 Combining {len(ordered_files)} source files in dependency order:")

    # Concatenate files
    concatenate_files(ordered_files, output_file)

    # Format with ruff
    run_ruff_format(output_file)

    # Report results
    with open(output_file, encoding="utf-8") as f:
        content = f.read()

    print("\n" + "=" * 70)
    print("✅ Build complete!")
    print("=" * 70)
    print(f"📍 Output: {output_file}")
    print(f"📏 Lines:  {len(content.splitlines()):,}")
    print(f"💾 Size:   {len(content):,} bytes")
    print("=" * 70)


if __name__ == "__main__":
    main()