annotate.py - comment on files and save comments for lines

annotate.py

#!/usr/bin/env python3
"""
annotate.py - Navigate a markdown plan file and record comments per line.

Comments are saved to comments.json (same directory as the file) on quit.

Usage: annotate.py <file>

Keys:
  j / DOWN    move down
  k / UP      move up
  g           go to top
  G           go to bottom
  c           add/edit comment on current line
  d           delete comment on current line
  q           quit and save comments.json
  Q           quit without saving
"""

import curses
import sys
import os
import json
import subprocess
import tempfile


def load_lines(path):
    with open(path) as f:
        return f.read().splitlines()


def load_comments(json_path):
    """Returns dict[int, str] keyed by 1-based line number."""
    if not os.path.exists(json_path):
        return {}
    with open(json_path) as f:
        data = json.load(f)
    return {entry["line"]: entry["comment"] for entry in data}


def save_comments(json_path, comments):
    data = [{"line": k, "comment": v} for k, v in sorted(comments.items())]
    with open(json_path, "w") as f:
        json.dump(data, f, indent=2)


def open_editor(prefill=""):
    """Open $EDITOR (fallback vim) on a temp file. Returns text or None if cancelled/empty."""
    editor = os.environ.get("EDITOR", "vim")
    with tempfile.NamedTemporaryFile(mode="w", suffix=".md", delete=False) as f:
        f.write(prefill)
        tmp = f.name
    try:
        subprocess.call([editor, tmp])
        with open(tmp) as f:
            text = f.read().strip()
        return text if text else None
    finally:
        os.unlink(tmp)


def input_line(stdscr, prompt, prefill=""):
    """Single-line input at the bottom of the screen. Returns text or None if cancelled."""
    h, w = stdscr.getmaxyx()
    row = h - 1

    curses.curs_set(1)
    buf = list(prefill)
    cursor = len(buf)

    def redraw():
        stdscr.move(row, 0)
        stdscr.clrtoeol()
        stdscr.addstr(row, 0, prompt, curses.A_REVERSE)
        px = len(prompt)
        avail = w - px - 1
        display = "".join(buf)
        if len(display) > avail:
            display = display[-avail:]
        stdscr.addstr(row, px, display)
        stdscr.move(row, px + min(cursor, avail))
        stdscr.refresh()

    redraw()

    result = None
    while True:
        try:
            ch = stdscr.get_wch()
        except curses.error:
            continue

        code = ord(ch) if isinstance(ch, str) else ch

        if code in (10, 13):
            result = "".join(buf)
            break
        elif code == 27:
            break
        elif code in (curses.KEY_BACKSPACE, 127, 8):
            if cursor > 0:
                buf.pop(cursor - 1)
                cursor -= 1
        elif code == curses.KEY_LEFT:
            cursor = max(0, cursor - 1)
        elif code == curses.KEY_RIGHT:
            cursor = min(len(buf), cursor + 1)
        elif code == curses.KEY_HOME:
            cursor = 0
        elif code == curses.KEY_END:
            cursor = len(buf)
        elif isinstance(ch, str) and ch.isprintable():
            buf.insert(cursor, ch)
            cursor += 1

        redraw()

    curses.curs_set(0)
    return result


def main(stdscr, path):
    curses.curs_set(0)
    curses.use_default_colors()
    curses.init_pair(1, curses.COLOR_CYAN, -1)            # lines with comments
    curses.init_pair(2, curses.COLOR_YELLOW, -1)          # line numbers
    curses.init_pair(3, curses.COLOR_BLACK, curses.COLOR_WHITE)  # selected
    curses.init_pair(4, curses.COLOR_BLACK, curses.COLOR_CYAN)   # selected + comment

    json_path = os.path.join(os.path.dirname(os.path.abspath(path)), "comments.json")
    lines = load_lines(path)
    comments = load_comments(json_path)  # dict[int, str], 1-based
    cursor = 0
    offset = 0
    dirty = False

    def adjust_scroll(view_h):
        nonlocal offset
        if cursor < offset:
            offset = cursor
        elif cursor >= offset + view_h:
            offset = cursor - view_h + 1

    def draw(message=""):
        stdscr.erase()
        h, w = stdscr.getmaxyx()
        view_h = h - 2
        adjust_scroll(view_h)
        num_w = len(str(len(lines))) + 1

        for row in range(view_h):
            idx = offset + row
            if idx >= len(lines):
                break
            lineno = idx + 1  # 1-based
            has_comment = lineno in comments
            is_cur = idx == cursor

            if is_cur:
                line_attr = curses.color_pair(4) if has_comment else curses.color_pair(3)
                num_attr = line_attr
            else:
                line_attr = curses.color_pair(1) if has_comment else curses.A_NORMAL
                num_attr = curses.color_pair(2)

            num_str = f"{lineno:>{num_w}} "
            stdscr.addstr(row, 0, num_str, num_attr)

            # Show comment preview inline after line content
            avail = w - num_w - 2
            content = lines[idx]
            if has_comment:
                preview = f"  // {comments[lineno]}"
                content_trunc = content[:max(0, avail - len(preview))]
                display = content_trunc + preview
            else:
                display = content[:avail]

            try:
                stdscr.addstr(row, num_w + 1, display[:avail], line_attr)
            except curses.error:
                pass

        # Status bar
        modified = " [modified]" if dirty else ""
        status = f" {os.path.basename(path)}{modified}  {cursor+1}/{len(lines)}  | c:comment  C:editor  d:delete  q:save+quit  Q:quit"
        stdscr.addstr(h - 2, 0, status[:w-1].ljust(w - 1), curses.A_REVERSE)

        if message:
            try:
                stdscr.addstr(h - 1, 0, message[:w-1], curses.A_BOLD)
            except curses.error:
                pass

        stdscr.refresh()

    message = ""
    while True:
        cursor = max(0, min(len(lines) - 1, cursor))
        draw(message)
        message = ""

        try:
            ch = stdscr.get_wch()
        except curses.error:
            continue

        if ch in ("j", curses.KEY_DOWN):
            cursor += 1
        elif ch in ("k", curses.KEY_UP):
            cursor -= 1
        elif ch == "g":
            cursor = 0
        elif ch == "G":
            cursor = len(lines) - 1
        elif ch == "C":
            lineno = cursor + 1
            prefill = comments.get(lineno, "")
            curses.endwin()
            text = open_editor(prefill=prefill)
            stdscr = curses.initscr()
            curses.curs_set(0)
            if text is not None:
                comments[lineno] = text
                dirty = True
                message = "Comment saved."
            else:
                message = "Empty — not saved." if not prefill else "Unchanged."
        elif ch == "c":
            lineno = cursor + 1
            prefill = comments.get(lineno, "")
            text = input_line(stdscr, f"Comment line {lineno}: ", prefill=prefill)
            if text is not None:
                if text.strip():
                    comments[lineno] = text
                    dirty = True
                    message = "Comment saved."
                else:
                    message = "Empty — not saved."
            else:
                message = "Cancelled."
        elif ch == "d":
            lineno = cursor + 1
            if lineno in comments:
                del comments[lineno]
                dirty = True
                message = "Comment deleted."
            else:
                message = "No comment on this line."
        elif ch == "q":
            if dirty:
                save_comments(json_path, comments)
            break
        elif ch == "Q":
            break


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(f"Usage: {sys.argv[0]} <file>")
        sys.exit(1)
    path = sys.argv[1]
    if not os.path.exists(path):
        print(f"File not found: {path}")
        sys.exit(1)
    curses.wrapper(main, path)