A CLI tool to generate Gantt charts from simulation run CSVs.

gantt.py

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.14"
# dependencies = [
#   "pandas",
#   "matplotlib",
#   "cyclopts",
# ]
# ///

import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.patches import Patch
from cyclopts import App
from pathlib import Path

app = App(help="A CLI tool to generate Gantt charts from simulation run CSVs.")

@app.default
def main(
    input_file: Path, 
    output_file: Path = Path("gantt_chart.png"),
    title: str = "Projected Simulation Schedule"
):
    """
    Reads a simulation CSV and produces a Gantt chart swimlaned by Allocated Resource.

    Parameters
    ----------
    input_file
        Path to the CSV file (e.g. BroadbandSim_Projected - Runs.csv).
    output_file
        Path where the resulting PNG should be saved.
    title
        The title displayed at the top of the chart.
    """
    
    if not input_file.exists():
        print(f"Error: File {input_file} not found.")
        return

    df = pd.read_csv(input_file)

    df['Projected Start Date'] = pd.to_datetime(df['Projected Start Date'], errors='coerce')
    df['Projected End Date'] = pd.to_datetime(df['Projected End Date'], errors='coerce')

    required_cols = ['Projected Start Date', 'Projected End Date', 'Allocated Resource', 'Run Tag']
    df = df.dropna(subset=required_cols)

    if df.empty:
        print("No valid scheduled entries found in the CSV.")
        return

    df = df.sort_values(by=['Allocated Resource', 'Projected Start Date'], ascending=[True, True])
    df['task_y_index'] = range(len(df))

    fig, ax = plt.subplots(figsize=(14, 8))
    
    researchers = df['Researcher'].unique()
    colours = plt.get_cmap('Set2', len(researchers))
    res_colour_map = {res: colours(i) for i, res in enumerate(researchers)}

    for _, row in df.iterrows():
        start = row['Projected Start Date']
        end = row['Projected End Date']
        duration = (end - start).days
        
        duration = max(duration, 0.5)
        
        ax.barh(row['task_y_index'], duration, left=start, height=0.6, 
                color=res_colour_map[row['Researcher']], 
                edgecolor='black', alpha=0.9)
        
        # Center-aligned text label for the Run Tag
        ax.text(start + pd.Timedelta(days=duration/2), row['task_y_index'], 
                row['Run Tag'], va='center', ha='center', 
                color='black', fontsize=9, fontweight='bold', clip_on=True)

    legend_elements = [Patch(facecolor=colour, edgecolor='k', label=researcher) for researcher, colour in res_colour_map.items()]

    ax.legend(handles=legend_elements)
    
    yticks = []
    yticklabels = []
    resources = df['Allocated Resource'].unique()
    for res in resources:
        res_indices = df[df['Allocated Resource'] == res]['task_y_index']
        yticks.append(res_indices.mean())
        yticklabels.append(res)
        
        ax.axhline(res_indices.max() + 0.5, color='black', linewidth=1, alpha=0.2)

    ax.set_yticks(yticks)
    ax.set_yticklabels(yticklabels)
    ax.invert_yaxis()

    ax.xaxis.set_major_locator(mdates.WeekdayLocator(interval=1))
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%b %d'))
    plt.xticks(rotation=45)

    ax.set_title(title, fontsize=16, pad=20)
    ax.set_xlabel('Timeline')
    ax.set_ylabel('Compute Resource')
    plt.grid(axis='x', linestyle='--', alpha=0.5)
    plt.tight_layout()

    plt.savefig(output_file)

if __name__ == "__main__":
    app()