Ipython extension for my qol improvements.

qol.py

"""IPython extension for quality of life improvements.

1. Shows execution time for every cell in blue (ANSI color in terminal)
2. Saves HTML representation of objects with _repr_html_ to ~/.local/ipython-html/
   and maintains an index.html with melange light theme for viewing
3. Configures plotly to use the melange color palette
"""

import os
import json
import time
from datetime import datetime
from IPython.core.getipython import get_ipython

# Store for timing data
_cell_start_time = None
_html_counter = 0


# ANSI color codes
BLUE = "\033[34m"
RESET = "\033[0m"


# Color palettes
palettes = {
    "melange": {
        # Core colors from melange light
        "bg": "#F1F1F1",
        "fg": "#54433A",
        "fg_dim": "#7D6658",
        "border": "#D9D3CE",
        "hover": "#E9E1DB",
        "black": "#E9E1DB",
        # Primary accent colors
        "red": "#C77B8B",
        "blue": "#7892BD",
        "green": "#6E9B72",
        "yellow": "#BC5C00",
        "cyan": "#739797",
        "magenta": "#BE79BB",
        # Bright variants for data visualization
        "bright_red": "#BF0021",
        "bright_blue": "#465AA4",
        "bright_green": "#3A684A",
        "bright_yellow": "#A06D00",
        "bright_cyan": "#3D6568",
        "bright_magenta": "#904180",
        # Dark backgrounds for contrast areas
        "dark_red": "#F1DEDF",
        "dark_blue": "#E0E2E8",
        "dark_green": "#D0E9D1",
        "dark_yellow": "#CCA478",
        "dark_cyan": "#CDE8E7",
        "dark_magenta": "#E8E0E8",
    },
    "melange_dark": {
        # Core colors from melange dark
        "bg": "#2A2522",
        "fg": "#ECE1D7",
        "fg_dim": "#A38A78",
        "border": "#4A443F",
        "hover": "#3A3530",
        "black": "#2A2522",
        # Primary accent colors (same as light)
        "red": "#C77B8B",
        "blue": "#7892BD",
        "green": "#6E9B72",
        "yellow": "#BC5C00",
        "cyan": "#739797",
        "magenta": "#BE79BB",
        # Bright variants
        "bright_red": "#BF0021",
        "bright_blue": "#465AA4",
        "bright_green": "#3A684A",
        "bright_yellow": "#A06D00",
        "bright_cyan": "#3D6568",
        "bright_magenta": "#904180",
        # Light backgrounds for contrast
        "dark_red": "#4A3A3C",
        "dark_blue": "#3A4454",
        "dark_green": "#3A4A3C",
        "dark_yellow": "#4A4030",
        "dark_cyan": "#3A4A4A",
        "dark_magenta": "#4A3A48",
    },
    "opencode": {
        # Core - Clean minimal light theme
        "bg": "#fafafa",
        "fg": "#1a1a1a",
        "fg_dim": "#666666",
        "border": "#e5e5e5",
        "hover": "#f5f5f5",
        "surface": "#ffffff",
        # Accents
        "accent": "#0066cc",
        "accent_hover": "#0052a3",
        # Data visualization colors
        "red": "#dc2626",
        "blue": "#2563eb",
        "green": "#16a34a",
        "yellow": "#ca8a04",
        "cyan": "#0891b2",
        "magenta": "#c026d3",
        # Darker variants for contrast
        "dark_red": "#fef2f2",
        "dark_blue": "#eff6ff",
        "dark_green": "#f0fdf4",
        "dark_yellow": "#fefce8",
        "dark_cyan": "#ecfeff",
        "dark_magenta": "#fdf4ff",
    },
}

# Active palette
active_palette = "opencode"


def _get_palette_colors():
    """Get the active palette color dict."""
    return palettes[active_palette]


def _ensure_html_directory():
    """Ensure the HTML storage directory exists and create static index files if missing."""
    html_dir = os.path.expanduser("~/.local/ipython-html/")
    os.makedirs(html_dir, exist_ok=True)

    # Create static index files if they don't exist
    _ensure_index_files(html_dir)

    return html_dir


def _ensure_index_files(html_dir):
    """Create static index.html and initial index.json if they don't exist."""
    index_html_path = os.path.join(html_dir, "index.html")
    index_json_path = os.path.join(html_dir, "index.json")

    if not os.path.exists(index_html_path):
        with open(index_html_path, "w", encoding="utf-8") as f:
            f.write(_generate_index_html_content())

    if not os.path.exists(index_json_path):
        with open(index_json_path, "w", encoding="utf-8") as f:
            json.dump({"files": [], "updated": datetime.now().isoformat()}, f)


def _get_html_files(html_dir):
    """Get list of HTML files (excluding index.html, index.json) sorted by modification time."""
    files = []
    for f in os.listdir(html_dir):
        if f.endswith(".html") and f not in ("index.html",):
            filepath = os.path.join(html_dir, f)
            files.append((os.path.getmtime(filepath), f))
    files.sort(reverse=True)  # Most recent first
    return [f for _, f in files]


def _update_index_json(html_dir):
    """Update only the index.json file with current list of HTML files."""
    files = _get_html_files(html_dir)
    index_data = {"files": files, "updated": datetime.now().isoformat()}
    index_path = os.path.join(html_dir, "index.json")
    with open(index_path, "w", encoding="utf-8") as f:
        json.dump(index_data, f)


def _generate_index_html_content():
    """Generate the index.html content with melange theme."""
    c = _get_palette_colors()

    return f"""<!DOCTYPE html>
<html>
<head>
    <title>outputs</title>
    <meta charset="utf-8">
    <style>
        /* Clean minimal theme - opencode palette */
        :root {{
            --bg: {c["bg"]};
            --fg: {c["fg"]};
            --fg-dim: {c["fg_dim"]};
            --border: {c["border"]};
            --hover: {c["hover"]};
            --accent: {c["accent"]};
            --surface: {c["surface"]};
            
            --iframe-height: 400px;
        }}
        
        * {{
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }}
        
        html {{
            margin: 0;
            padding: 0;
        }}
        
        body {{
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
            font-size: 13px;
            line-height: 1.4;
            color: var(--fg);
            background: var(--bg);
            margin: 0;
            padding: 4px 8px;
        }}
        
        .controls {{
            display: flex;
            align-items: center;
            gap: 12px;
            margin: 0;
            padding: 6px 0;
            border-bottom: 1px solid var(--border);
            flex-wrap: wrap;
        }}
        
        .k-control {{
            display: flex;
            align-items: center;
            gap: 4px;
            color: var(--fg-dim);
            font-size: 12px;
        }}
        
        .k-control button {{
            background: {c["surface"]};
            border: 1px solid var(--border);
            color: var(--fg);
            width: 22px;
            height: 22px;
            cursor: pointer;
            font-size: 12px;
            line-height: 1;
            transition: all 0.15s;
        }}
        
        .k-control button:hover {{
            background: var(--hover);
            border-color: var(--fg-dim);
        }}
        
        .k-control button:disabled {{
            opacity: 0.3;
            cursor: default;
        }}
        
        .k-value {{
            min-width: 16px;
            text-align: center;
            color: var(--fg);
            font-weight: 500;
        }}
        
        .height-control {{
            display: flex;
            align-items: center;
            gap: 6px;
            color: var(--fg-dim);
            font-size: 12px;
        }}
        
        .height-control input[type="range"] {{
            width: 80px;
            cursor: pointer;
            -webkit-appearance: none;
            appearance: none;
            height: 3px;
            background: var(--border);
            outline: none;
        }}
        
        .height-control input[type="range"]::-webkit-slider-thumb {{
            -webkit-appearance: none;
            appearance: none;
            width: 12px;
            height: 12px;
            background: {c["surface"]};
            border: 2px solid var(--fg-dim);
            cursor: pointer;
        }}
        
        .height-control input[type="range"]::-moz-range-thumb {{
            width: 12px;
            height: 12px;
            background: {c["surface"]};
            border: 2px solid var(--fg-dim);
            cursor: pointer;
        }}
        
        .height-value {{
            min-width: 40px;
            color: var(--fg);
            font-size: 12px;
        }}
        
        .status {{
            color: var(--fg-dim);
            font-size: 12px;
            margin-left: auto;
        }}
        
        .items {{
            display: flex;
            flex-direction: column;
            gap: 8px;
        }}
        
        .item {{
            border: 1px solid var(--border);
            background: {c["surface"]};
            overflow: hidden;
        }}
        
        .item-header {{
            display: flex;
            justify-content: space-between;
            align-items: center;
            padding: 6px 10px;
            background: var(--hover);
            border-bottom: 1px solid var(--border);
            font-size: 11px;
            color: var(--fg-dim);
        }}
        
        .item-header a {{
            color: var(--fg);
            text-decoration: none;
            font-weight: 500;
        }}
        
        .item-header a:hover {{
            color: var(--accent);
        }}
        
        .item-content iframe {{
            width: 100%;
            border: none;
            display: block;
            height: var(--iframe-height);
        }}
        
        .empty {{
            color: var(--fg-dim);
            font-style: italic;
            padding: 20px;
            text-align: center;
        }}
    </style>
</head>
<body>
    <div class="controls">
        <div class="k-control">
            <span>show</span>
            <button id="minus">-</button>
            <span class="k-value" id="k">3</span>
            <button id="plus">+</button>
        </div>
        <div class="height-control">
            <span>height</span>
            <input type="range" id="heightSlider" min="100" max="1000" value="400" step="50">
            <span class="height-value" id="heightValue">400px</span>
        </div>
        <span class="status" id="status">loading...</span>
    </div>
    
    <div id="items" class="items"></div>

    <script>
        let k = 3;
        const minK = 1, maxK = 20;
        let currentFiles = [];
        let iframeHeight = 400;
        
        // Load from storage
        try {{
            const savedK = localStorage.getItem('k');
            if (savedK) k = parseInt(savedK);
            const savedHeight = localStorage.getItem('iframeHeight');
            if (savedHeight) iframeHeight = parseInt(savedHeight);
        }} catch(e) {{}}
        
        document.getElementById('k').textContent = k;
        document.getElementById('heightSlider').value = iframeHeight;
        document.getElementById('heightValue').textContent = iframeHeight + 'px';
        document.documentElement.style.setProperty('--iframe-height', iframeHeight + 'px');
        
        function updateButtons() {{
            document.getElementById('minus').disabled = k <= minK;
            document.getElementById('plus').disabled = k >= maxK;
        }}
        
        document.getElementById('minus').onclick = () => {{
            if (k > minK) {{ k--; saveK(); updateVisibility(); }}
        }};
        
        document.getElementById('plus').onclick = () => {{
            if (k < maxK) {{ k++; saveK(); updateVisibility(); }}
        }};
        
        document.getElementById('heightSlider').oninput = (e) => {{
            iframeHeight = parseInt(e.target.value);
            document.getElementById('heightValue').textContent = iframeHeight + 'px';
            document.documentElement.style.setProperty('--iframe-height', iframeHeight + 'px');
            try {{ localStorage.setItem('iframeHeight', iframeHeight); }} catch(e) {{}}
        }};
        
        function saveK() {{
            document.getElementById('k').textContent = k;
            try {{ localStorage.setItem('k', k); }} catch(e) {{}}
            updateButtons();
        }}
        
        function formatTime(filename) {{
            const parts = filename.split('_');
            if (parts.length < 2) return filename;
            const date = parts[0];
            const time = parts[1];
            if (!date || !time || date.length !== 8 || time.length !== 6) return filename;
            return `${{date.slice(0,4)}}-${{date.slice(4,6)}}-${{date.slice(6,8)}} ${{time.slice(0,2)}}:${{time.slice(2,4)}}:${{time.slice(4,6)}}`;
        }}
        
        // Track rendered items: filename -> DOM element
        const renderedItems = new Map();
        
        function createItemElement(f) {{
            const div = document.createElement('div');
            div.className = 'item';
            div.dataset.filename = f;
            div.innerHTML = `
                <div class="item-header">
                    <a href="${{f}}" target="_blank">${{f}}</a>
                    <span>${{formatTime(f)}}</span>
                </div>
                <div class="item-content"><iframe src="${{f}}"></iframe></div>
            `;
            return div;
        }}
        
        function updateVisibility() {{
            // Get items in current DOM order (which matches files array order)
            const container = document.getElementById('items');
            const items = Array.from(container.querySelectorAll('.item'));
            items.forEach((element, index) => {{
                element.style.display = index < k ? 'block' : 'none';
            }});
        }}
        
        async function render(files) {{
            const container = document.getElementById('items');
            
            // Handle empty state
            if (files.length === 0) {{
                // Clear all rendered items
                renderedItems.forEach((element) => element.remove());
                renderedItems.clear();
                container.innerHTML = '<div class="empty">no outputs</div>';
                return;
            }}
            
            // Remove empty message if it exists
            const emptyMsg = container.querySelector('.empty');
            if (emptyMsg) emptyMsg.remove();
            
            // Remove items that are no longer in files list
            renderedItems.forEach((element, filename) => {{
                if (!files.includes(filename)) {{
                    element.remove();
                    renderedItems.delete(filename);
                }}
            }});
            
            // Add or reorder items
            files.forEach((f, index) => {{
                let element = renderedItems.get(f);
                
                if (!element) {{
                    // Create new element
                    element = createItemElement(f);
                    renderedItems.set(f, element);
                }}
                
                // Ensure correct order in DOM
                if (container.children[index] !== element) {{
                    if (index < container.children.length) {{
                        container.insertBefore(element, container.children[index]);
                    }} else {{
                        container.appendChild(element);
                    }}
                }}
            }});
            
            // Update visibility based on k
            updateVisibility();
        }}
        
        async function checkIndex() {{
            try {{
                const r = await fetch('index.json?t=' + Date.now());
                const data = await r.json();
                const newFiles = data.files || [];
                
                const changed = newFiles.length !== currentFiles.length ||
                    newFiles.some((f, i) => f !== currentFiles[i]);
                
                if (changed) {{
                    currentFiles = newFiles;
                    document.getElementById('status').textContent = 
                        currentFiles.length + ' items';
                    await render(currentFiles);
                }}
            }} catch(e) {{
                document.getElementById('status').textContent = 'error';
            }}
        }}
        
        async function load() {{
            await checkIndex();
        }}
        
        updateButtons();
        load();
        
        setInterval(checkIndex, 1000);
    </script>
</body>
</html>"""


def _wrap_html_with_css(html_content):
    """Wrap HTML content with CSS styling for proper table rendering.

    If content is already a full HTML document (contains <html>), return as-is.
    Otherwise, wrap it with proper HTML structure and CSS.
    """
    # Check if already a full HTML document
    if "<html" in html_content.lower():
        return html_content

    # Use active palette colors
    c = _get_palette_colors()
    return f"""<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <style>
        body {{
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
            font-size: 12px;
            line-height: 1.4;
            color: {c["fg"]};
            background: {c["bg"]};
            margin: 0;
            padding: 8px;
        }}
        
        table {{
            border-collapse: collapse;
            width: 100%;
            font-size: 12px;
            background: {c["bg"]};
        }}
        
        th, td {{
            border: 1px solid {c["border"]};
            padding: 4px 8px;
            text-align: left;
        }}
        
        th {{
            background: {c["hover"]};
            font-weight: 500;
            color: {c["fg"]};
        }}
        
        tr:nth-child(even) {{
            background: {c["hover"]};
        }}
        
        thead tr:nth-child(2) {{
            font-style: italic !important;
        }}
        
        small {{
            color: {c["fg_dim"]};
        }}
        
        .dataframe {{
            border: 1px solid {c["border"]};
        }}
        
        .dataframe thead th {{
            background: {c["hover"]};
        }}
    </style>
</head>
<body>
{html_content}
</body>
</html>"""


def configure_plotly_theme():
    """Configure plotly to use the active color palette and suppress __repr__."""
    try:
        import plotly.io as pio
        import plotly.graph_objects as go

        # Monkey patch Figure to have empty __repr__
        original_repr = go.Figure.__repr__

        def empty_repr(self):
            return ""

        go.Figure.__repr__ = empty_repr

        c = _get_palette_colors()

        # Define the color sequence for plotly using palette colors
        color_sequence = [
            c.get("blue", "#2563eb"),
            c.get("red", "#dc2626"),
            c.get("green", "#16a34a"),
            c.get("yellow", "#ca8a04"),
            c.get("cyan", "#0891b2"),
            c.get("magenta", "#c026d3"),
            c.get("bright_blue", "#3b82f6"),
            c.get("bright_red", "#ef4444"),
            c.get("bright_green", "#22c55e"),
            c.get("bright_yellow", "#eab308"),
        ]

        # Create a custom template
        template = go.layout.Template()

        # Layout defaults
        template.layout = {
            "paper_bgcolor": c["bg"],
            "plot_bgcolor": c["bg"],
            "font": {"color": c["fg"], "family": "sans-serif"},
            "title": {"font": {"color": c["fg"]}},
            "legend": {"font": {"color": c["fg"]}, "bgcolor": "rgba(0,0,0,0)"},
            "xaxis": {
                "gridcolor": c["border"],
                "linecolor": c["border"],
                "tickcolor": c["border"],
                "title": {"font": {"color": c["fg"]}},
                "tickfont": {"color": c["fg_dim"]},
            },
            "yaxis": {
                "gridcolor": c["border"],
                "linecolor": c["border"],
                "tickcolor": c["border"],
                "title": {"font": {"color": c["fg"]}},
                "tickfont": {"color": c["fg_dim"]},
            },
            "colorway": color_sequence,
        }

        # Register the template
        pio.templates["melange"] = template
        pio.templates.default = "melange"

        # Disable auto-rendering
        pio.renderers.default = None

    except ImportError:
        pass  # plotly not installed


def _pre_run_cell(info):
    """Hook called before cell execution starts."""
    global _cell_start_time
    _cell_start_time = time.time()


def _post_run_cell(info):
    """Hook called after cell execution completes."""
    global _cell_start_time, _html_counter

    # Calculate execution time
    if _cell_start_time is not None:
        elapsed = time.time() - _cell_start_time
        # Display time in blue using ANSI escape codes
        print(f"{BLUE}[{elapsed:.3f} s]{RESET}")
        _cell_start_time = None

    # Check if there's a result with _repr_html_
    result = info.result
    if result is not None and hasattr(result, "_repr_html_"):
        try:
            html_content = result._repr_html_()
            if html_content is not None:
                html_dir = _ensure_html_directory()
                _html_counter += 1
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"{timestamp}_{_html_counter}.html"
                filepath = os.path.join(html_dir, filename)

                with open(filepath, "w", encoding="utf-8") as f:
                    f.write(_wrap_html_with_css(html_content))
                print(f"{BLUE}Saved HTML to {filepath}{RESET}")

                # Update only index.json (index.html is static)
                _update_index_json(html_dir)
        except Exception:
            # Silently ignore errors in _repr_html_ processing
            pass


def load_ipython_extension(ipython):
    """Load the extension."""
    # Configure plotly theme on load
    configure_plotly_theme()

    ipython.events.register("pre_run_cell", _pre_run_cell)
    ipython.events.register("post_run_cell", _post_run_cell)


def unload_ipython_extension(ipython):
    """Unload the extension."""
    ipython.events.unregister("pre_run_cell", _pre_run_cell)
    ipython.events.unregister("post_run_cell", _post_run_cell)