Add init --sync and --prune for incremental monitor setup

`init` planned every file on every run and refused if any of them
existed, so attaching a new monitor could not be configured without
--force — which reallocated global workspace numbers by i3's current
output order and renumbered the monitors already in use.

- allocate_global_starts() keeps the block each known output already
  holds and gives new outputs the lowest free block; blocks of inactive
  outputs stay reserved.
- read_allocations()/parse_allocation() read those blocks back out of
  the generated files (inverse of render_workspaces_for_output).
- --sync writes only the missing files, --prune drops configs for
  outputs that are gone and frees their numbers.
- Files without the auto-generated header are never written or
  deleted, --force included.

Co-Authored-By: Claude <noreply@anthropic.com>

pi3/smart_workspace.py

@@ -33,7 +33,7 @@ import pprint
 import re
 import sys
 from dataclasses import dataclass
-from typing import Iterable
+from typing import Iterable, Mapping
 
 import i3ipc
 
@@ -451,6 +451,101 @@ def render_bindings(count: int, mod: str = "$mod") -> str:
     return "".join(out)
 
 
+@dataclass(frozen=True)
+class Allocation:
+    """A block of global workspace numbers already handed to an output.
+
+    Attributes:
+        output: The output name the block belongs to.
+        start: First global workspace number in the block.
+        count: How many consecutive numbers the block covers.
+    """
+
+    output: str
+    start: int
+    count: int
+
+
+_ALLOC_OUTPUT_RE = re.compile(r"^\s*set\s+\$out_\w+\s+(\S+)\s*$", re.MULTILINE)
+_ALLOC_WS_RE = re.compile(r"^\s*set\s+\$ws_\w+\s+(\d+):\d+\s*$", re.MULTILINE)
+
+
+def parse_allocation(config_text: str) -> Allocation | None:
+    """Read back the output and global-number block from a generated file.
+
+    The inverse of :func:`render_workspaces_for_output`, used by
+    ``init --sync`` to keep the numbers a monitor already has instead of
+    reallocating them from scratch.
+
+    Args:
+        config_text: Contents of a ``workspaces.conf`` previously written
+            by ``init``.
+
+    Returns:
+        The :class:`Allocation` described by the file, or ``None`` if the
+        text doesn't contain both an ``$out_…`` alias and at least one
+        ``$ws_…`` alias (i.e. it isn't one of our generated files).
+    """
+    output_match = _ALLOC_OUTPUT_RE.search(config_text)
+    numbers = [int(m.group(1)) for m in _ALLOC_WS_RE.finditer(config_text)]
+    if output_match is None or not numbers:
+        return None
+    return Allocation(output_match.group(1), min(numbers), len(numbers))
+
+
+def allocate_global_starts(
+    outputs: Iterable[str],
+    existing: Mapping[str, Allocation],
+    count: int,
+) -> dict[str, int]:
+    """Decide each output's first global workspace number.
+
+    Outputs that already hold a block of the right size keep it, so
+    plugging in a new monitor never renumbers the ones you already use.
+    New outputs get the lowest block of ``count`` free numbers.
+
+    Blocks belonging to outputs in ``existing`` that aren't in ``outputs``
+    (an unplugged monitor) stay reserved — a screen that is merely
+    switched off must not lose its numbers to a newcomer. Drop those
+    entries from ``existing`` (what ``--prune`` does) to free them.
+
+    Args:
+        outputs: Names of the outputs to allocate for, in i3's order.
+        existing: Already-allocated blocks, keyed by output name.
+        count: Workspaces per output.
+
+    Returns:
+        ``{output_name: global_start}`` covering every name in ``outputs``.
+    """
+    if count < 1:
+        raise ValueError("count must be >= 1")
+
+    occupied: set[int] = set()
+    for alloc in existing.values():
+        occupied.update(range(alloc.start, alloc.start + alloc.count))
+
+    starts: dict[str, int] = {}
+    pending: list[str] = []
+    for name in outputs:
+        alloc = existing.get(name)
+        if alloc is not None and alloc.count == count:
+            starts[name] = alloc.start
+        else:
+            if alloc is not None:
+                # The requested block size changed (``-n``), so the old
+                # block can't be kept — release it and reallocate below.
+                occupied.difference_update(range(alloc.start, alloc.start + alloc.count))
+            pending.append(name)
+
+    for name in pending:
+        start = 1
+        while not occupied.isdisjoint(range(start, start + count)):
+            start += 1
+        starts[name] = start
+        occupied.update(range(start, start + count))
+    return starts
+
+
 @dataclass(frozen=True)
 class InitArgs:
     """Parsed CLI arguments for the ``init`` subcommand."""
@@ -459,6 +554,8 @@ class InitArgs:
     config_dir: pathlib.Path
     dry_run: bool
     force: bool
+    sync: bool
+    prune: bool
     debug: bool
 
 
@@ -466,19 +563,32 @@ def _planned_files(
     config_dir: pathlib.Path,
     outputs: list[tuple[str, Rect]],
     count: int,
+    starts: Mapping[str, int] | None = None,
 ) -> list[tuple[pathlib.Path, str]]:
     """Compute the (path, content) pairs that ``init`` would write.
 
     Pure function: no filesystem or i3 access. Each output gets a folder
     ``<config_dir>/<output_name>/workspaces.conf``; the shared keybindings
-    go in ``<config_dir>/bindings.conf``. Output order is preserved from
-    ``outputs`` so global workspace numbers are deterministic.
+    go in ``<config_dir>/bindings.conf``.
+
+    Args:
+        config_dir: Directory the generated files live in.
+        outputs: Active outputs, in i3's order.
+        count: Workspaces per output.
+        starts: Optional ``{output: global_start}`` map, typically from
+            :func:`allocate_global_starts`. When omitted, each output gets
+            a contiguous block based on its position in ``outputs`` —
+            deterministic, but it renumbers everything if the output order
+            changes.
     """
     files: list[tuple[pathlib.Path, str]] = [
         (config_dir / "bindings.conf", render_bindings(count)),
     ]
     for index, (output_name, _) in enumerate(outputs):
-        global_start = index * count + 1
+        if starts is not None:
+            global_start = starts[output_name]
+        else:
+            global_start = index * count + 1
         files.append(
             (
                 config_dir / output_name / "workspaces.conf",
@@ -488,15 +598,73 @@ def _planned_files(
     return files
 
 
+def _is_generated(path: pathlib.Path) -> bool:
+    """Return whether ``path`` carries our auto-generation header.
+
+    Anything else is treated as hand-written and never overwritten or
+    deleted, not even with ``--force``.
+    """
+    try:
+        return path.read_text().startswith(_AUTOGEN_HEADER)
+    except OSError:
+        return False
+
+
+def read_allocations(config_dir: pathlib.Path) -> dict[str, Allocation]:
+    """Read the global-number blocks already handed out under ``config_dir``.
+
+    Scans ``<config_dir>/*/workspaces.conf`` and parses each generated
+    file with :func:`parse_allocation`. Hand-edited or unreadable files
+    are skipped, so their numbers are *not* considered reserved.
+
+    Args:
+        config_dir: The (possibly non-existent) init config directory.
+
+    Returns:
+        ``{output_name: Allocation}``; empty if nothing has been generated
+        yet.
+    """
+    allocations: dict[str, Allocation] = {}
+    for path in sorted(config_dir.glob("*/workspaces.conf")):
+        if not _is_generated(path):
+            continue
+        try:
+            alloc = parse_allocation(path.read_text())
+        except OSError:
+            continue
+        if alloc is not None:
+            allocations[alloc.output] = alloc
+    return allocations
+
+
+def _prune_paths(config_dir: pathlib.Path, stale: Iterable[str]) -> list[pathlib.Path]:
+    """Return the generated ``workspaces.conf`` files of no-longer-active outputs."""
+    paths = []
+    for name in stale:
+        path = config_dir / name / "workspaces.conf"
+        if path.exists() and _is_generated(path):
+            paths.append(path)
+    return paths
+
+
 def run_init(args: InitArgs) -> None:
     """Generate per-output i3 workspace configs for the currently-active outputs.
 
+    Three write modes:
+
+    * default — write everything, but refuse if any target file exists.
+    * ``--sync`` — write only the files that are missing (the new
+      monitor), leaving existing outputs' numbers untouched.
+    * ``--force`` — regenerate everything from scratch, which reallocates
+      global workspace numbers by i3's current output order.
+
     Args:
         args: Parsed init-subcommand arguments.
 
     Raises:
-        RuntimeError: If no active outputs are detected, or if files
-            already exist and ``--force`` was not passed.
+        RuntimeError: If no active outputs are detected, if files already
+            exist and neither ``--sync`` nor ``--force`` was passed, or if
+            a target file exists but wasn't generated by us.
     """
     i3 = i3ipc.Connection()
     outputs = _active_outputs(i3)
@@ -504,27 +672,82 @@ def run_init(args: InitArgs) -> None:
         raise RuntimeError("no active outputs detected by i3")
 
     config_dir = args.config_dir.expanduser()
-    files = _planned_files(config_dir, outputs, args.count)
+    active_names = [name for name, _ in outputs]
+
+    allocations = read_allocations(config_dir)
+    stale = [name for name in allocations if name not in active_names]
+
+    if args.force:
+        # Explicit "start over": numbers follow i3's output order.
+        starts = None
+    else:
+        if args.prune:
+            for name in stale:
+                allocations.pop(name)
+        starts = allocate_global_starts(active_names, allocations, args.count)
+
+    files = _planned_files(config_dir, outputs, args.count, starts)
+    pruned = _prune_paths(config_dir, stale) if args.prune else []
+
+    if args.sync:
+        planned = [(path, content) for path, content in files if not path.exists()]
+    else:
+        planned = files
 
     if args.dry_run:
-        for path, content in files:
+        for path in pruned:
+            print(f"--- would remove {path} ---")
+        for path, content in planned:
             print(f"--- {path} ---")
             print(content, end="")
+        if not planned:
+            print("nothing to write; every active output already has a config")
         return
 
-    existing = [path for path, _ in files if path.exists()]
+    existing = [path for path, _ in planned if path.exists()]
     if existing and not args.force:
         listing = "\n  ".join(str(p) for p in existing)
         raise RuntimeError(
-            "refusing to overwrite existing files (pass --force to replace):\n  "
-            + listing
+            "refusing to overwrite existing files (pass --sync to only add "
+            "configs for new outputs, or --force to replace):\n  " + listing
         )
 
-    for path, content in files:
+    # Never clobber something the user wrote by hand — --force included.
+    hand_written = [path for path in existing if not _is_generated(path)]
+    if hand_written:
+        listing = "\n  ".join(str(p) for p in hand_written)
+        raise RuntimeError(
+            "refusing to overwrite files that were not generated by "
+            "`pi3-smart-workspace init` (move them aside first):\n  " + listing
+        )
+
+    for path, content in planned:
         path.parent.mkdir(parents=True, exist_ok=True)
         path.write_text(content)
         print(f"wrote {path}")
 
+    for path in pruned:
+        path.unlink()
+        print(f"removed {path}")
+        try:
+            path.parent.rmdir()
+        except OSError:
+            pass  # directory not empty — leave whatever else is in there
+
+    if not planned:
+        print("nothing to write; every active output already has a config")
+    if args.sync and stale and not args.prune:
+        listing = ", ".join(stale)
+        print(f"note: configs kept for inactive outputs: {listing} (--prune removes them)")
+    if args.sync:
+        bindings, bindings_content = files[0]
+        if bindings.exists() and _is_generated(bindings):
+            if bindings.read_text() != bindings_content:
+                print(
+                    f"note: {bindings} differs from what -n {args.count} would "
+                    "generate; pass --force to regenerate it"
+                )
+
     print()
     print("Add the following to your i3 config (typically ~/.config/i3/config):")
     print(f"  include {config_dir}/bindings.conf")
@@ -550,9 +773,22 @@ def _init_main(argv: list[str]) -> None:
         "--dry-run", action="store_true",
         help="Print generated configs to stdout instead of writing them.",
     )
-    parser.add_argument(
+    mode = parser.add_mutually_exclusive_group()
+    mode.add_argument(
+        "--sync", action="store_true",
+        help="Only write configs for outputs that don't have one yet, keeping "
+             "the workspace numbers of the outputs that do. Use this after "
+             "plugging in a new monitor.",
+    )
+    mode.add_argument(
         "--force", action="store_true",
-        help="Overwrite existing files in the config directory.",
+        help="Regenerate every file from scratch. This reallocates global "
+             "workspace numbers by i3's current output order.",
+    )
+    parser.add_argument(
+        "--prune", action="store_true",
+        help="Also delete generated configs for outputs that are no longer "
+             "active, freeing their workspace numbers.",
     )
     parser.add_argument(
         "-d", "--debug", action="store_true",
@@ -565,6 +801,8 @@ def _init_main(argv: list[str]) -> None:
         config_dir=parsed.config_dir,
         dry_run=parsed.dry_run,
         force=parsed.force,
+        sync=parsed.sync,
+        prune=parsed.prune,
         debug=parsed.debug,
     )