Small script for plotting graphs from text logs

regex_plot.py

#!/usr/bin/env python3

import re
import sys
import argparse
import datetime

import matplotlib.pyplot as plt


def parse_x(s, date_format):
    try:
        return float(s)
    except ValueError:
        if date_format is None:
            raise

    return datetime.datetime.strptime(s, date_format)


def regex_type_helper(s):
    return (repr(s), re.compile(s))


def assignment_type_helper(s):
    return (
        s,
        re.compile(
            re.escape(s) + r"\s*=\s*([-+]?[0-9]*\.?[0-9]+(?:[eE][-+]?[0-9]+)?)"
        )
    )


add_yaxis_flag = object()


parser = argparse.ArgumentParser(description="Regex plotter")
parser.add_argument(
    "--regex", "-r",
    action="append",
    dest="plots",
    metavar="RE",
    type=regex_type_helper,
    help="Regex capturing values to plot. Must capture one or two values."
)
parser.add_argument(
    "--assignment", "-a",
    action="append",
    dest="plots",
    metavar="ID",
    type=assignment_type_helper,
    help="Shortcut for plotting expressions in form of 'ID = value'."
)
parser.add_argument(
    "--add-yaxis",
    action="append_const",
    const=add_yaxis_flag,
    dest="plots",
    help="All following values will be ploted on a new secondary Y axis"
)
parser.add_argument(
    "--x-regex", "-x",
    default=None,
    metavar="RE",
    type=regex_type_helper,
    help="Regex for obtaining X value of a data point. Must capture single value."
)
parser.add_argument(
    "--date-format",
    help="Format used for parsing dates"
)
parser.add_argument(
    "--vertical-mark",
    action="append",
    dest="marks",
    metavar="RE",
    type=regex_type_helper,
    help="Place a vertical mark on the plot. Must capture zero, one or two values."
)
parser.add_argument(
    "--start-regex",
    type=regex_type_helper,
    default=None,
    help="Don't process any data until this regex matches"
)
parser.add_argument(
    "file",
    type=argparse.FileType("r"),
    nargs="*",
    default=[sys.stdin]
)

args = parser.parse_args()

data_x = {}
data_y = {}
marks_x = {}

started = args.start_regex is None

mark_prev_values = {}

for file in args.file:
    for line in file:
        line = line[:-1]

        if not started:
            match = args.start_regex[1].search(line)
            if match:
                started = True
            else:
                continue

        if args.x_regex is not None:
            match = args.x_regex[1].search(line)
            if not match:
                line_x_coord = None
            else:
                groups = match.groups()
                if len(groups) != 1:
                    raise ValueError("X regex must capture exactly one value")

                line_x_coord = parse_x(groups[0], args.date_format)

        print(line)

        for plot in args.plots:
            try:
                name, regex = plot
            except TypeError:
                pass

            match = re.search(regex, line)
            if not match:
                continue

            data_x_list = data_x.setdefault(name, [])
            data_y_list = data_y.setdefault(name, [])

            groups = match.groups()

            if len(groups) == 1:
                if line_x_coord is None:
                    x_coord = len(data_x_list)
                else:
                    x_coord = line_x_coord
            elif len(groups) == 2:
                x_coord = parse_x(groups[0], args.date_format)
            else:
                raise ValueError("Regex must capture 1 or 2 values")

            y_coord = float(groups[-1])

            print(f"Plot {name}: ({x_coord}, {y_coord})")

            data_x_list.append(x_coord)
            data_y_list.append(y_coord)

        for mark in args.marks:
            name, regex = mark

            match = re.search(regex, line)
            if not match:
                continue

            mark_x_list = marks_x.setdefault(name, [])

            groups = match.groups()

            if len(groups) == 0 and line_x_coord is not None:
                x_coord = line_x_coord
                value = object()  # Always unique
            elif len(groups) == 1:
                if line_x_coord is None:
                    x_coord = parse_x(groups[0], args.date_format)
                    value = object()  # Always unique
                else:
                    x_coord = line_x_coord
                    value = groups[0]
            elif len(groups) == 2:
                x_coord = parse_x(groups[0], args.date_format)
                value = groups[1]
            else:
                raise ValueError("Regex must capture 1 or 2 values")

            if value != mark_prev_values.get(name, None):
                mark_x_list.append(x_coord)
                print(f"Mark {name}: ({x_coord})")
            mark_prev_values[name] = value

fig, axes = plt.subplots()
axes.grid()
lines = []
axes_suffix = ""

fig.canvas.manager.set_window_title(', '.join(f.name for f in args.file))

for plot in args.plots:
    if plot is add_yaxis_flag:
        prop_cycler = axes._get_lines.prop_cycler
        axes = axes.twinx()
        axes._get_lines.prop_cycler = prop_cycler  # Don't restart the line style cycler
        axes_suffix += "→"
    else:
        name = plot[0]
        if name in data_x:
            n = name
            if len(axes_suffix):
                n += " (" + axes_suffix + ")"
            lines += axes.plot(data_x[name], data_y[name], label=n)
        else:
            print("No data for", name)

for mark in args.marks:
    name, regex = mark
    if name in marks_x:
        for x in marks_x[name]:
            axes.axvline(x, label=name)
    else:
        print("No data for", name)


axes.legend(lines, [ln.get_label() for ln in lines])

fig.tight_layout()
plt.show()