lean-split-cli.py

#!/usr/bin/env python3
"""CLI client for lean-split-tool module selection and flake generation."""

import argparse
import os
import re
import sys

SITE_URL = "https://meta-introspector.github.io/lean-split-tool"


def parse_dot_file(dot_path):
    """Parse modules.dot to extract module list"""
    modules = set()
    with open(dot_path, "r") as f:
        content = f.read()
    for m in re.finditer(r'"([^"]+)"\s*->|node\s+"([^"]+)"', content):
        modules.add(m.group(1) or m.group(2))
    return sorted(modules)


def search_modules(modules, term):
    """Search modules by term"""
    term = term.lower()
    return [m for m in modules if term in m.lower()]


def generate_nix_cmd(modules):
    """Generate nix build commands"""
    return (
        "nix build "
        + " \\\n  ".join(
            f'"github:meta-introspector/mathlib4?ref=feature/split&dir={m.replace(".", "/")}"'
            for m in modules
        )
        + " --no-write-lock-file"
    )


def generate_flake(modules, repo=None, branch=None):
    """Generate flake.nix content"""
    repo = repo or "meta-introspector/mathlib4"
    branch = branch or "feature/split"
    inputs = ['    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";']
    for m in modules:
        mod_name = m.replace(".", "_").replace("-", "_")
        inputs.append(
            f'    {mod_name}.url = "github:{repo}?ref={branch}&dir={m.replace(".", "/")}";'
        )

    return f"""{{
  description = "Generated mathlib module flake";
  inputs = {{
{chr(10).join(inputs)}
  }};
  outputs = {{ self, nixpkgs, ... }}@inputs:
    let
      system = "x86_64-linux";
      pkgs = nixpkgs.legacyPackages.${{system}};
    in {{
      packages.${{system}}.default = pkgs.stdenv.mkDerivation {{
        pname = "mathlib-modules";
        version = "0.1.0";
        src = ./.;
      }};
    }};
}}"""


def find_lean_project_root(start_dir):
    """Find the root directory of a Lean project (contains lean-toolchain or lakefile)."""
    current = start_dir
    while True:
        if os.path.exists(os.path.join(current, "lean-toolchain")):
            return current
        if os.path.exists(os.path.join(current, "lakefile.lean")) or os.path.exists(os.path.join(current, "lakefile.toml")):
            return current
        parent = os.path.dirname(current)
        if parent == current or not parent:
            return start_dir  # Return original if no root found
        current = parent


def get_source_directory(project_root):
    """Determine if project uses Mathlib/ or RequestProject/ structure."""
    mathlib_path = os.path.join(project_root, "Mathlib")
    request_project_path = os.path.join(project_root, "RequestProject")
    
    if os.path.exists(request_project_path):
        return "RequestProject"
    elif os.path.exists(mathlib_path):
        return "Mathlib"
    else:
        # Fallback: check which directory has .lean files
        for root, dirs, files in os.walk(project_root):
            for f in files:
                if f.endswith('.lean'):
                    # Check if in Mathlib or RequestProject subdirectory
                    if 'Mathlib' in root.split(os.sep):
                        return "Mathlib"
                    elif 'RequestProject' in root.split(os.sep):
                        return "RequestProject"
        return "Mathlib"  # Default assumption


def build_lean_module_list(project_root, source_dir):
    """Build a list of Lean modules from the project structure."""
    modules = []
    
    # Walk through the source directory to find .lean files
    for root, dirs, files in os.walk(source_dir):
        for f in files:
            if f.endswith('.lean') and f != 'All.lean':
                full_path = os.path.join(root, f)
                rel_path = os.path.relpath(full_path, project_root)
                # Convert path to module name
                module_name = rel_path.replace('.lean', '').replace(os.sep, '.')
                modules.append(module_name)
    
    return sorted(set(modules))  # Deduplicate and sort


def parse_lean_module_file(lean_file_path):
    """Parse a .lean file to extract its module name."""
    try:
        with open(lean_file_path, 'r') as f:
            for line in f:
                if line.strip().startswith('module '):
                    return line.strip().split(' ', 1)[1].rstrip('\n')
    except Exception:
        pass
    # Fallback: try to extract from path structure
    rel_path = os.path.relpath(lean_file_path, os.path.dirname(lean_file_path))
    return rel_path.replace('.lean', '').replace(os.sep, '.')


def main():
    parser = argparse.ArgumentParser(
        description="Lean Split CLI - Select and build Lean modules"
    )
    parser.add_argument(
        "command",
        choices=["list", "search", "nix", "flake", "build", "split"],
        help="Command to run",
    )
    parser.add_argument(
        "query", nargs="?", help="Search term, module list, or path (for split)"
    )
    parser.add_argument(
        "--dot", default="site/modules.dot", help="Path to modules.dot file"
    )
    parser.add_argument("--output", "-o", help="Output file for flake generation")
    parser.add_argument(
        "--repo",
        default="meta-introspector/mathlib4",
        help="Custom mathlib fork (owner/repo format)",
    )
    parser.add_argument(
        "--branch",
        default="feature/split",
        help="Branch for modular flakes",
    )

    args = parser.parse_args()

    if os.path.exists(args.dot):
        modules = parse_dot_file(args.dot)
    else:
        try:
            import urllib.request

            with urllib.request.urlopen(f"{SITE_URL}/modules.dot") as resp:
                content = resp.read().decode()
            modules = re.findall(r'"([^"]+)"', content)
        except Exception as e:
            print(f"Error fetching modules: {e}")
            sys.exit(1)

    if args.command == "list":
        for m in modules[:50]:
            print(m)
        print(f"... ({len(modules)} total modules)")

    elif args.command == "search":
        if not args.query:
            print("Error: search term required")
            sys.exit(1)
        results = search_modules(modules, args.query)[:50]
        for m in results:
            print(m)
        print(f"Found {len(results)} modules matching '{args.query}'")

    elif args.command == "nix":
        if not args.query:
            print("Error: module list required")
            sys.exit(1)
        selected = args.query.replace(",", " ").split()
        valid = [
            m for m in selected if m in modules or any(m in mod for mod in modules)
        ]
        if not valid:
            print(f"No valid modules found in: {args.query}")
            print("Use 'search <term>' to find modules")
            sys.exit(1)
        print(generate_nix_cmd(valid))

    elif args.command == "flake":
        if not args.query:
            print("Error: module list required")
            sys.exit(1)
        selected = args.query.replace(",", " ").split()
        valid = [
            m for m in selected if m in modules or any(m in mod for mod in modules)
        ]
        content = generate_flake(valid, args.repo, args.branch)
        if args.output:
            with open(args.output, "w") as f:
                f.write(content)
            print(f"Written to {args.output}")
        else:
            print(content)

    elif args.command == "split":
        if not args.query:
            print("Error: path to mathlib source required")
            print(
                "Usage: nix run .#lean-split-cli split /path/to/mathlib --repo owner/repo"
                "\n       nix run .#lean-split-cli split /path/to/project --repo owner/repo [--source-dir RequestProject]"
            )
            sys.exit(1)
        
        project_root = find_lean_project_root(args.query)
        source_structure = get_source_directory(project_root)
        
        if source_structure == "RequestProject":
            # For RequestProject structure, provide explicit module discovery
            print(f"⚠️  Detected RequestProject structure in {args.query}")
            print(f"   ℹ️  Source directory: {source_structure}")
            
            # Find the RequestProject directory
            request_project_path = os.path.join(project_root, "RequestProject")
            if not os.path.exists(request_project_path):
                print(f"Error: RequestProject directory not found at {request_project_path}")
                sys.exit(1)
            
            # Build list of modules from RequestProject
            modules_list = build_lean_module_list(project_root, request_project_path)
            print(f"Found {len(modules_list)} Lean modules in RequestProject")
            
            if not modules_list:
                print("No .lean files found in RequestProject directory")
                sys.exit(1)
            
            # For now, we'll use the split script approach but need to handle differently
            print(f"Splitting project at: {args.query}")
            print(f"Target: {args.repo}#{args.branch}")
        else:
            # For Mathlib structure, use existing logic
            source_dir = os.path.join(args.query, "Mathlib") if source_structure == "Mathlib" else args.query
            print(f"Splitting mathlib at: {args.query}")
            print(f"Source directory: {source_dir}")
        
        print(f"Target: {args.repo}#{args.branch}")


if __name__ == "__main__":
    main()