analyze_profiles.py

# File: analyze_profiles.py
import logging
import pstats
from pathlib import Path
from pstats import SortKey
from typing import Annotated

import typer

# Configure logging for the script itself
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
logger = logging.getLogger(__name__)

app = typer.Typer(
    help="Analyzes cProfile output files (.prof) generated by AlphaTriangle workers."
)

# Define annotations for Typer arguments/options
ProfilePathArg = Annotated[
    Path,
    typer.Argument(
        ...,  # Ellipsis makes it a required argument
        help="Path to the .prof file to analyze.",
        exists=True,
        file_okay=True,
        dir_okay=False,
        readable=True,
        resolve_path=True,
    ),
]

# --- Corrected Option Definition (Default in Function Signature) ---
# Remove default from Option(), add it to the function parameter below
NumLinesOption = Annotated[
    int,
    typer.Option("--lines", "-n", help="Number of lines to show."),
]
# --- End Correction ---


@app.command()
def analyze(
    profile_path: ProfilePathArg,  # Use the annotation
    num_lines: NumLinesOption = 30,  # Assign default value here
):
    """
    Loads a .prof file and prints statistics sorted by cumulative and total time.
    """
    logger.info(f"Analyzing profile: {profile_path}")

    try:
        # Ensure profile_path is converted to string for pstats
        p = pstats.Stats(str(profile_path))
    except Exception as e:
        logger.error(f"Error loading profile stats from {profile_path}: {e}")
        # Add 'from e' to preserve original exception context
        raise typer.Exit(code=1) from e

    # Remove directory paths for cleaner output
    p.strip_dirs()

    print("\n" + "=" * 30)
    print(f"  Top {num_lines} Functions by Cumulative Time")
    print("=" * 30)
    try:
        p.sort_stats(SortKey.CUMULATIVE).print_stats(num_lines)
    except Exception as e:
        logger.error(f"Error sorting/printing by cumulative time: {e}")

    print("\n" + "=" * 30)
    print(f"  Top {num_lines} Functions by Total Internal Time")
    print("=" * 30)
    try:
        p.sort_stats(SortKey.TIME).print_stats(num_lines)  # SortKey.TIME is Tottime
    except Exception as e:
        logger.error(f"Error sorting/printing by total time: {e}")

    logger.info(f"Finished analyzing {profile_path}")


if __name__ == "__main__":
    app()