simplify_game_log.py

#!/usr/bin/env python
# -*- coding: utf-8 0*0

import json
import os
import multiprocessing as mp
import tqdm

FILESDIRECTORY = 'data'
OUTPUTFILE = 'games.txt'
winner_map = {'none': 0, 'player1': 1, 'player2':2}
pbar = None

def simplify(game_log_str):
    try:
        game_log = json.loads(game_log_str)
        winner = winner_map[game_log['settings']['winnerplayer']]
        states_list = game_log['states'][1::2]

        moves = []
        previous_state = ','.join(['0'] * 81)
        for state in states_list:
            moves.append(find_move(previous_state, state['field']))
            previous_state = state['field']
        return {'winner': winner, 'moves': moves}
    except ValueError:
        pass


def find_move(previous_state, current_state):
    previous_state_list = previous_state.split(',')
    current_state_list = current_state.split(',')

    for i in range(len(previous_state_list)):
        if previous_state_list[i] == '0' and current_state_list[i] != '0':
            return str(i)
    return '-1'


def evaluate(filename, queue):
    with open(FILESDIRECTORY + '/' + filename) as f:
        game_log_str = f.read()
        game = simplify(game_log_str)
        if game and len(game['moves']) > 14:
            message = str(game['winner']) + ';' + (','.join(game['moves']))
            queue.put(message)


def queue_listener(queue):
    with open(OUTPUTFILE, 'wb') as f:
        while True:
            message = queue.get()
            if message == 'done':
                break
            pbar.update()
            f.write(message + '\n')

if __name__ == '__main__':
    file_list = os.listdir(FILESDIRECTORY)
    pbar = tqdm.tqdm(total=len(file_list))

    manager = mp.Manager()
    queue = manager.Queue()
    pool = mp.Pool(mp.cpu_count() * 2)

    listener = pool.apply_async(queue_listener, (queue, ))

    tasks = []
    for filename in file_list:
        task = pool.apply_async(evaluate, (filename, queue))
        tasks.append(task)

    for task in tasks:
        task.get()

    queue.put('done')
    pbar.close()
    pool.close()