RoboSanta/json2csv.py at master · iliakur/RoboSanta · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
"""Turns a collection of stats JSON files into a CSV.

The point of separating this that collecting all the stats is potentially costly,
so it's better to do that once, get as much as possible and store that,
then extract specific information from them into CSV files taylored to a specific purpose.
"""
from collections import OrderedDict, namedtuple
import csv
import re

from toolz import curry
import click

from fileIO import experiment_item_name, iterate_dir, read_json

ExperimentData = namedtuple("ExperimentData", "name summary")


def parse_experiment_data(stats_path: str):
    stats = read_json(stats_path)
    return ExperimentData(experiment_item_name(stats_path), stats['summary'])


def get_grid_size(data: ExperimentData):
    m = re.search("_n(\d+)_", data.name)
    return m.group(1) if m else 0


def _get_time(data: ExperimentData, key: str):
    return data.summary['times'][key]


def get_solve_time(data: ExperimentData):
    return _get_time(data, 'solve')


def get_total_time(data: ExperimentData):
    return _get_time(data, 'total')


def get_solved(data: ExperimentData):
    return data.summary['result'] == 1.0 and get_solve_time(data) > 0.0


def get_makespan(data: ExperimentData):
    return data.summary['costs'][0]


# yapf: disable
ENTRIES = OrderedDict([
    ('name', lambda data: data.name),
    ('grid_size', get_grid_size),
    ('solved', get_solved),
    ('total_time', get_total_time),
    ('solve_time', get_solve_time),
    ('plan_makespan', get_makespan)])
# yapf: enable


@curry
def create_row(fields: OrderedDict, data: ExperimentData):
    return dict((field_name, field_getter(data)) for field_name, field_getter in fields.items())

@click.command()
@click.argument("stats_dir")
@click.argument("output_path")
def main(stats_dir, output_path):
    with open(output_path, mode='w') as fh:
        writer = csv.DictWriter(fh, fieldnames=ENTRIES.keys())
        writer.writeheader()
        writer.writerows(
            map(create_row(ENTRIES), map(parse_experiment_data, iterate_dir(stats_dir))))


if __name__ == '__main__':
    main()