Shortens each edge label in a LJA DOT file to just the first line

shorten_edge_labels.py

#! /usr/bin/env python
#
# Shortens edge labels in a DOT file output by LJA to just show the first line
# and then a count of how many other lines are omitted. (If an edge's label
# spans exactly one or two lines, then the entire label is preserved.)
#
# USAGE:
# ./shorten_edge_labels.py in.dot out.dot

import sys

if len(sys.argv) != 3:
    raise ValueError("Run this command as ./shorten_edge_labels.py in.dot out.dot")

INFILE = sys.argv[1]
OUTFILE = sys.argv[2]

with open(INFILE, "r") as df:
    with open(OUTFILE, "w") as of:
        for line in df:
            if "->" in line:
                # shorten line label to just first two lines
                halves = line.split('[label="')
                label = halves[1].split('"')[0]
                # remove empty lines ("") caused by e.g. a trailing
                # \\n, or something like \\n\\n in the middle of the label
                label_lines = [ll for ll in label.split("\\n") if len(ll) > 0]
                if len(label_lines) == 1:
                    suffix = ""
                elif len(label_lines) == 2:
                    # no sense saying "1 more line" when we can just
                    # include the line itself here
                    suffix = f"\\n{label_lines[-1]}"
                else:
                    # the - 1 accounts for how we show the first line
                    suffix = f"\\n{len(label_lines) - 1:,} more lines"
                of.write(f'{halves[0]}[label="{label_lines[0]}{suffix}"]\n')
            else:
                of.write(line)