POC parse JSON to CSV

csv.py

#!/usr/bin/env python

import unittest

import argparse
from argparse import ArgumentParser
import json
import os
import sys

def get_columns(benchmarks):
    """
    Gather all possible column names from the JSON
    """
    cols = set()
    for bm in benchmarks:
        cols.update(bm.keys())
    return sorted(cols)


def main():
    parser = ArgumentParser(description='convert json benchmark output to csv')
    parser.add_argument(
        '-i', '--input',
        nargs='?',
        type=argparse.FileType('r'),
        default=sys.stdin,
        dest='input',
        help='JSON file containing output from a benchmark run',
    )

    args, unknown_args = parser.parse_known_args()
    if args.input is None:
        parser.print_help()
        exit(1)
    if len(unknown_args) != 0:
        print('unknown args: ', unknown_args)
        parser.print_help()
        exit(1)

    input_json = json.load(args.input)
    benchmarks = input_json['benchmarks']

    cols = get_columns(benchmarks)
    print(','.join(cols))

    for bm in benchmarks:
        vals = [str(bm[col]) if col in bm else '' for col in cols]
        print(','.join(vals))

    args.input.close()

### TESTS ###
class TestGetColumns(unittest.TestCase):
    def setUp(self):
        test_inputs = os.path.join(
            os.path.dirname(os.path.realpath(__file__)),
            '..', 'gbench', 'Inputs')
        test_input1 = os.path.join(test_inputs, 'test1_run1.json')
        self.test_json = json.load(open(test_input1))

    def test_get_columns_input1(self):
        self.assertEqual(
            get_columns(self.test_json['benchmarks']),
            [
                'aggregate_name', 'big_o', 'cpu_coefficient', 'cpu_time', 'iterations',
                'label', 'name', 'real_coefficient', 'real_time', 'rms','run_name',
                'run_type', 'time_unit'
            ])


if __name__ == '__main__':
    # unittest.main()
    main()