dir_walker.py

import json
import os
import pickle

from simulation.serialization import RESULTS_DIR

# Directory walker with custom functions invoked on each subdirectory (recursively).
# It can be utilized as results processor/fixing tool.

def walk(results_dir, fun):
    rootDir = results_dir
    for dirName, subdirList, fileList in os.walk(rootDir):
        print('Found directory: %s' % dirName)
        for fname in fileList:
            fun(dirName, fname)

def json2pickle(dirName, fname):
    if fname.endswith('.json'):
        to_convert = os.path.join(dirName, fname)
        with open(to_convert, 'r') as fh:
            loaded = json.load(fh)

            to_save = os.path.join(dirName, fname[:-4] + 'pickle')
            with open(to_save, 'wb') as fw:
                pickle.dump(loaded, fw)
                print('\tsaved pickle: %s' % to_save)
        os.remove(to_convert)


def clear_budget(dirName, fname):
    splitted = fname.split('.')
    budget = int(splitted[0])
    print('\t %s' % fname)
    if budget > 4500:
        to_remove = os.path.join(dirName, fname)
        print('\tRemoving: %s' % to_remove)
        os.remove(to_remove)

if __name__ == '__main__':
    walk(RESULTS_DIR, json2pickle)