csv2gnuplot

#!/usr/bin/env python3
"""
  Input a CSV or similarly delimited file, output it to gnuplot.

  The input files expected here have multiple data sets split across rows and columns, in a way that makes
  the files easy to create in a spreadsheet or a logging data collector.  That information is parsed
  into a data set format built with Python dictionary and list structures before going into the plotting core.

  Special support is in place for files that come out of the timed-os-stats utility.
  The original, more generic csv2gnuplot function is still in here as well.  Sample data for each
  format is only in the doc comments of each function below so far.

  Very rough code right now filled with TODO items.  This program aims to be a standalone tool that needs
  only very basic Python and gnuplot.  The same job could be done with more advanced Python data tools,
  like storing into a Pandas DataFrame and then using Matplotlib to render it.  Some of the users of
  this code will not have those libraries available.  A long-term goal is to provide a functionally similar
  program to this one that's based on a proper scientific Python stack instead.

  Copyright 2011-2023, Greg Smith.
  Part of the pgbench-tools set
"""

import csv
import optparse
import sys
import os
import time
from dateutil import parser

import logging
dblevel=logging.INFO
logging.basicConfig(level=dblevel)
logger = logging.getLogger(__name__)
logger.setLevel(dblevel)

def parse_options():
    optparser = optparse.OptionParser(
        usage="usage: %prog [options] [input_file] [output_file]",
        version="1.0",
        conflict_handler="resolve")

    optparser.add_option('-i', '--input', dest="input_file", default=None,
                         help="File to read from.  Defaults to standard input, except when overwritten by arguments after the regular options.")

    optparser.add_option('-o', '--output', dest="output_file", default="graph.png",
                         help="File to write output to.  Defaults to graph.png")

    optparser.add_option('-t', '--title', dest="title", default="PostgreSQL pgbench SELECT-Only Rate",
                         help="Title of graph")

    optparser.add_option('-x', '--xlabel', dest="xlabel", default="Clients",
                         help="X Axis Label")

    optparser.add_option('-y', '--ylabel', dest="ylabel", default="Transactions per second",
                         help="Y Axis Label")

    optparser.add_option('-s', '--signal', dest="header_signal", default=None,
                         help="Header signal")

    optparser.add_option('-d', '--data', dest="data_type", default=None,
                         help="Type of data file.  options are 'table', 'iostat' and 'vmstat'")

    optparser.add_option('--disks', dest="disk_list", default='',
                         help="Comma separated list of disks to include in iostat format")

    options, args = optparser.parse_args()

    if len(args) > 1:
        options.output_file = args[1]
    if len(args) > 0:
        options.input_file = args[0]

    return options, args

def try_pandas(s,fn):
    # TODO Make this a command line option
    want_pandas = False
    if not want_pandas:
        return

    try:
        import pandas as pd
    except:
        return

    df=pd.DataFrame(s)
    df.to_json(fn)
    logger.debug("pandas saved to '%s'",fn)

# TODO Refactor the common gnuplot setup tasks in table_gnuplotter and ds_gnuplotter into gnuplot wrapper functions

def table_gnuplotter(rows, header=True, destination="graph", title="", xlabel=None,ylabel=None):
    """
    Outputs a list of lists into gnuplot

        - rows - list of lists
        - header - if True the first row is treated as a table header
        - destination - file to write the table to.  Defaults to graph.png

    This data format I'm thinking of as a 3D grid input, where one of the dimensions
    can be converted to data sets for easy plotting.  The real distinction here is
    that the first row is a header key describing one axis.  Here's an example:

    "Drive",12,24,48,96
    "DC-S3700",3682,6285,7581,9787
    "With BBWC",1951,2171,2345,2298
    "7200RPM",573,908,1306,1616

    That header line is saying that "Drive" is the first column, and the then second through fifth columns
    will have data entries for the values 12, 24, 48, and 96.  Each line here will turn into
    its own plot on the graph.

    """

    try:
        gp = os.popen('gnuplot', 'w')

        gp.write("set output '" + destination + ".png';\n")
        gp.write("set terminal pngcairo size 640,480 enhanced font 'sans,10';\n")

        gp.write("set grid xtics ytics;\n")
        gp.write("set key right bottom;\n")

        gp.write("set xlabel '%s';\n" % xlabel)
        gp.write("set ylabel '%s';\n" % ylabel)
        gp.write("set title '%s'; " % title)

        # Sometimes needed if auto-detection doesn't work right
        if (False):
            gp.write("set yrange [0:30];\n")

        gp.write("plot ")

        # Each line here is a new gnuplot data set file.  Iterate over the
        # list once to describe all of them.
        header = True
        header_text = None
        for rownum, row in enumerate(rows):
            if header:
                header_text = row
                header = False
                continue
            key = row[0]
            delimiter = ""
            if rownum < (len(rows) - 1):
                delimiter = ", "
            gp.write("'-' using 1:2 title '%s' with lines%s" % (key, delimiter))
        gp.write("\n")

        if header_text is not None:
            logger.debug("header text is" + str(header_text))

        # Now output the file data
        header = True
        for row in rows:
            if header:
                header = False
                continue

            # Iterate over the columns.  If not empty, print header for
            # that position, then the value.
            for column, value in enumerate(row):
                logger.debug("column=%s value=%s" % (column, value))
                if column == 0:
                    continue
                elif value is None or value.strip() == "":
                    continue
                else:
                    gp.write("%s %s\n" % (header_text[column], value))

            # Write gnuplot file delimiter
            gp.write("e\n")

        gp.close()
    except:
        print ("Error while spawning gnuplot")

def table_csv2gnuplot(in_file, out_file, graph_title, xlabel, ylabel):
    in_lines = []
    reader = csv.reader(in_file)
    for row in reader:
        in_lines.append(row)
    table_gnuplotter(in_lines, header=True, destination=out_file, title=graph_title,xlabel=xlabel,ylabel=ylabel)


def ds_gnuplotter(set_container, destination, title="", tsdata=True, ycolumn=1, ylabel="Rate"):
    """
    Outputs data sets (dictionary of dictionary of lists) into a gnuplot graph

        - sets - dictionary of dictionary (with data set as key) of lists
        - header - if True the first row is treated as a table header
        - destination - file to write the table to.  Defaults to standard out.

    This expects that the first column in all data sets is a timestamp used for
    the X axis values.  The time is in UNIX epoch format.
    """
    try:
        gp = os.popen('gnuplot', 'w')

        # TODO Make the input destination actually get used by the output file here
        gp.write("set output '" + destination + "';\n")
        gp.write("set terminal pngcairo size 640,480 enhanced font 'sans,10';\n")
        gp.write("set title '%s'; " % title)

        if tsdata:
            gp.write("set xdata time;\n")
            gp.write("set timefmt \"%s\";\n")
            gp.write("set format x \"%H:%M:%S\"\n")
            # When x axis label are in HH:MM::SS format, they will often overlap in common
            # gnuplot version+font combinations.  The problem seems related to how errors
            # like "using internal non-scalable font" will show up regularly.
            # This hard-coded workaround is hackish, but it works well enough for
            # the fixed 640x480 resolution of this program.
            gp.write("set xtics 10;\n")

        gp.write("set grid xtics ytics;\n")
        gp.write("set key right bottom;\n")
        gp.write("set xlabel 'Time'; set ylabel '%s'; " % ylabel)

        # Sometimes needed if auto-detection doesn't work right
        if False:
            gp.write("set yrange [0:30];\n")

        gp.write("plot ")

        # To keep the structure orthogonal, the outer dictionary here will have a
        # key name like "iostat" or "vmstat".  This code only handles one
        # value there, so just dereference that with the first entry found.
        logger.debug(set_container)
        for outer in set_container.keys():
            sets=set_container[outer]
            break

        # Each key here is a new gnuplot data set file.  Iterate over the
        # list once to describe all of them.
        logger.debug(sets)

        # TODO Convert key building to operate directly on dictionary
        # Building a list of the keys in this dictionary is done only to
        # detect when the last one is being processed, so that no "," is added.
        # There should be a way to do that directly in dictionary form.
        # enumerate works on dictionaries, but the result doesn't seem iterable.
        keys = []
        for key, rows in sets.items():
            keys.append(key)
        logger.debug("Data set list: %s\n" % keys)

        # List the data sets.  They all have to be defined before any of their data is written.
        # TODO Is it possible to define a set in a gnuplot input file and then immediately write its data?
        # This would be simpler if that's the case.
        for index, key in enumerate(keys):
            delimiter = ""
            if index < (len(keys) - 1):
                delimiter = ", "
            gp.write("'-' using 1:2 title '%s' with lines%s" % (key, delimiter))
        gp.write("\n")

        # Now output the file data
        for k, v in sets.items():
            logger.debug("v=%s", v)
            for column, data in enumerate(v):
                logger.debug("column=%s data=%s" % (column, data))
                if len(data) < (ycolumn + 1):
                    continue
                x = data[0]
                # gnuplot expects epoch timestamp data will have no fractional part.
                # Cope with that given data at this point is still in string form.
                if tsdata:
                    x = int(float(x))
                y = data[ycolumn]
                if y is None or y.strip() == "":
                    continue
                else:
                    logger.debug("(x,y)=%s %s\n" % (x, y))
                    gp.write("%s %s\n" % (x, y))

            # Write gnuplot file delimiter
            gp.write("e\n")
        gp.close()
    except:
        print ("Error while spawning gnuplot")

def value_parse(value):
    """
    Convert string value into a float.
    Device columns will have the device name, skip those and other junk.
    On Darwin some values have suffixes like this:
        188921K
    Those are mainly seen in the summary values given after each header.
    Not sure how high those can go; supporting K, M, and G.
    """

    multiplier=1
    if value.endswith("K"):  multiplier=1024
    if value.endswith("M"):  multiplier=1024*1024
    if value.endswith("G"):  multiplier=1024*1024*1024
    value=value.rstrip("KMG")  # You down with KMG?
    if multiplier>1:
        logger.debug('multiplier %s %s',value,multiplier)
    try:
        return float(value) * multiplier
    except:
        return None

def statdata2gnuplot(in_file, output_file, graph_title, filt_text, header_signal=None, setlabel=None,fake_signals=0):
    """
    Provides a generic way to derive labels that takes advantage of how some common input
    types--vmstat, iostat--are formatted.

    For iostat, the key is finding a line with "Device" and splitting that via a delimiter:

        Device:         rrqm/s   wrqm/s     r/s     w/s    rMB/s    wMB/s avgrq-sz avgqu-sz   await r_await w_await  svctm  %util

    For vmstat, looking for "cache" might do, then splitting this:

        r  b   swpd   free   buff  cache   si   so    bi    bo   in   cs us sy id wa st

    This assumes that the first 2 columns are a timestamp that's the X axis, which is the format
    output by the matching timed-os-stats program that feeds into this one

    Filters based on matching one of the entries in the filt_text

    header_signal is text that you'll find in a header line.  Useful for
    iostat and vmstat style data, where you can find out all of the column
    headers by parsing the right line.
    """

    search_for = None
    if filt_text is not None:
        search_for = filt_text.split(",")
    logger.debug("search_for: '%s' - filt_text '%s'" % (search_for,filt_text))

    if setlabel is None:
        print ("Error:  setlabel must be a string usable as an index.  It cannot be undefined.")
        return

    sets = {}
    csv_rows=[]
    header_labels = None
    prev_line = None
    signals = 0
    darwin_format=False
    gulp_next=False

    for row in in_file:
        columns = row.split()

        # Throw out Darwin lines like this right out
        # "Mach Virtual Memory Statistics: (page size of 4096 bytes)"
        if len(columns)>3:
            if columns[2]=="Mach" and columns[3]=="Virtual":
                darwin_format=True
                prev_line=row
                continue
        disks=None
        # Look for a header label row.
        # Stop processing once one is found, since programs like vmstat
        # republish the header every screen.
        if header_signal is not None:
            if row.find(header_signal) >= 0:
                signals=signals + 1
                if header_labels is None:
                    logger.debug("Previous line to header:  %s" % prev_line)
                    prev_columns=prev_line.split()
                    logger.debug("Previous line columns:  %s" % prev_columns)
                    header_labels = columns[2:]
                    logger.debug("Header labels:  %s" % header_labels)

                    # Darwin iostat:  previous line will have the disk device then "cpu"
                    # Use that sketchy set of facts to figure out the disk.
                    try:
                        logger.debug("previous columns index 3=%s" % prev_columns[3])
                        if prev_columns[3]=="cpu":
                            darwin_format=True
                            disk=prev_columns[2]
                            logger.debug("Disk %s" % disk)
                            header_labels[0]=disk+"_"+header_labels[0]
                            header_labels[1]=disk+"_"+header_labels[1]
                            header_labels[2]=disk+"_"+header_labels[2]

                        elif prev_columns[3].startswith("disk"):
                            # Handle two disk devices manually here.
                            # TODO Handle more than two
                            # TODO Reduce code duplication by handling 1 or 2
                            darwin_format=True

                            disk=prev_columns[2]
                            logger.debug("Disk %s" % disk)
                            header_labels[0]=disk+"_"+header_labels[0]
                            header_labels[1]=disk+"_"+header_labels[1]
                            header_labels[2]=disk+"_"+header_labels[2]

                            disk=prev_columns[3]
                            logger.debug("Disk %s" % disk)
                            header_labels[3]=disk+"_"+header_labels[3]
                            header_labels[4]=disk+"_"+header_labels[4]
                            header_labels[5]=disk+"_"+header_labels[5]

                    except:
                        logger.debug("Failed to process previous row")

                if header_labels is not None:
                    # Another sketchy Darwin detection.  Eventually this
                    # may need to just take Darwin format as an input
                    # parameter.
                    if "0fill" in header_labels:
                        darwin_format=True
                        gulp_next=True
                        prev_line=row
                        continue

        prev_line=row

        # Can't proceed if we haven't figured out which column to graph yet
        if header_labels is None:
            logger.debug("No header labels yet, skipping row %s" % row)
            continue

        if signals <= fake_signals:
            logger.debug("fake signal section, rejecting row %s" % row)
            # Fake signal on Darwin is a single line; processing anything to
            # this point means it's over
            if darwin_format:
                signals=signals + 1
            continue

        # Gobble Darwin vm_stat data republishing a summary after each header
        if gulp_next:
            gulp_next=False
            continue

        match_prefix=""
        matches = False
        if search_for is None:
            matches = True
            # When there isn't really a key to filter data sets, use an empty
            # string.  All useful lines will end up in that set.
            dataset = setlabel
        else:
            for s in search_for:
                if row.find(s) >= 0:
                    matches = True
                    dataset = str(s)
                    # Right now the only use case for this code is filtering
                    # disk_list iostat rows.  Without that, Linux iostat data
                    # will overflow with loopback and other unimportant disk
                    # traffic.  The disk entry that matched is then used
                    # as a prefix , giving labels like "sdb_rMB/s".
                    if len(s)>0:
                        # Find disk label after the "timestamp:" prefix 
                        dataset = (row[row.find(s):]).split()[0]
                        match_prefix=dataset + "_"
                    logger.debug("Searching for '%s' dataset '%s' using match_prefix '%s'",s,dataset,match_prefix)

        if not matches:
            logger.debug("no match - rejecting row %s" % row)
            logger.debug("search_for:  %s" % row)
            continue

        tstext = "%s %s" % (columns[0], columns[1])
        logger.debug(tstext)
        logger.debug(columns)

        # This doesn't save all the resolution on millisecond level timestamps.
        # But gnuplot don't handle them either though, so it's not a problem.
        # They get truncated to second resolution anyway by later code.
        ts = parser.parse(tstext)
        epoch = time.mktime(ts.timetuple())

        for pos, label in enumerate(header_labels):
            if len(columns) <= (pos + 2):
                logger.debug("No match for column %s in '%s'" % (label, row ))
                continue

            val = columns[pos + 2]
            logger.debug("Considering %s %s %s %s" % (dataset, label, epoch, val))

            float_val=value_parse(val)
            if float_val is None:
                logger.debug("Rejecting as not numeric:  '%s'",val)
                continue

            formatted = (str(epoch), str(float_val))

            if dataset not in sets:
                sets[str(dataset)] = {}
            if label not in sets[str(dataset)]:
                sets[str(dataset)][label] = []

            sets[str(dataset)][label].append(formatted)
            logger.debug("Appended %s %s %s %s" % (dataset, label, epoch, float_val))
            metric_row=[ts,float_val,match_prefix + label]
            csv_rows.append(metric_row)

# The CSV data has to be saved in an array and written at the end.
# TODO Consider opening the CSV early and writing to the CSV line at a time
    if len(csv_rows)>0:
        csvfn="%s.csv" % output_file
        with open(csvfn, 'w') as csvfile:
            fields = ['collected','value','metric']
            csvwriter = csv.writer(csvfile)
            csvwriter.writerow(fields)
            csvwriter.writerows(csv_rows)
            logger.info("metrics saved to '%s.csv'",output_file)

    logger.debug("Sets as they come out of the data set parsing")
    logger.debug(sets)
    if len(sets)>0:
        try_pandas(sets,output_file+'.json')
        ds_gnuplotter(sets, destination=output_file+'.png', title=graph_title)

def main():
    (options, args) = parse_options()

    if options.input_file is None:
        file_to_process = sys.stdin
    else:
        file_to_process = open(options.input_file, 'r')

    if options.output_file is None:
        out_file = "graph"
    else:
        out_file = options.output_file

    # TODO Convert remaining fixed values to parameter inputs:  filt_text, ylabel
    if options.data_type == "table":
        return table_csv2gnuplot(file_to_process,out_file,options.title,xlabel=options.xlabel,ylabel=options.ylabel)
    if options.data_type == "iostat":
        return statdata2gnuplot(in_file=file_to_process, output_file=out_file, graph_title=options.title,
                               filt_text=options.disk_list, 
                               header_signal=options.header_signal,
                               fake_signals=1,
                               setlabel="iostat")
    elif options.data_type == "vmstat":
        return statdata2gnuplot(in_file=file_to_process, output_file=out_file, graph_title=options.title,
                               filt_text=None, 
                               header_signal=options.header_signal,
                               setlabel="vmstat")
    else:
        print("Error:  unknown data file type %s" % options.data_type)

if __name__ == '__main__':
    main()