From b3c84c688425e346c2ea8f79dfacd35ac8420339 Mon Sep 17 00:00:00 2001 From: fishpm Date: Sun, 11 Sep 2016 15:53:32 +0200 Subject: [PATCH] Added RetrosheetTableSetter, db writing added to RetrosheetTableBuilder, small mods to RetrosheetParser and tests_etl --- etl/RetrosheetParser.py | 45 +++++- etl/RetrosheetTableBuilder.py | 283 ++++++++++++++++++++++------------ etl/RetrosheetTableSetter.py | 220 ++++++++++++++++++++++++++ tests/test_etl.py | 27 ++-- 4 files changed, 464 insertions(+), 111 deletions(-) create mode 100644 etl/RetrosheetTableSetter.py diff --git a/etl/RetrosheetParser.py b/etl/RetrosheetParser.py index 09fc771..8bd060f 100644 --- a/etl/RetrosheetParser.py +++ b/etl/RetrosheetParser.py @@ -5,10 +5,13 @@ ? Use python schematics for data validation """ +import re class RetrosheetLineParser(object): def __init__(self): + + # Each new line is parsed according to its first word. self.parser_map = { "id": self.parse_id , "version": self.parse_version @@ -25,7 +28,11 @@ def __init__(self): def type_parser(self, retro_list, parser_map): "Read the type and pass to the appropriate parser" + + # First word in line defines its type and how to be parsed. retro_type = retro_list[0] + + # Parse identifiable lines. if parser_map.has_key(retro_type): parser = self.parser_map[retro_type] type_parsed = parser(retro_list) @@ -37,6 +44,7 @@ def type_parser(self, retro_list, parser_map): def parse_id(self, retro_list): "Parse the id type" + game_id = retro_list[1] team = game_id[0:3] year = game_id[3:7] @@ -50,33 +58,49 @@ def parse_id(self, retro_list): , 'month': month , 'day': day , 'game_num': game_num + , 'type': 'id' } return id_instance def parse_version(self, retro_list): "Parse the version type" + # Indicated on retrosheet.org that the version record is obsolete. pass def parse_info(self, retro_list): "Parse the info type" - info_instance = {retro_list[1]: retro_list[2]} + + info_instance = {'type': 'info', retro_list[1]: retro_list[2]} return info_instance def parse_start(self, retro_list): "Parse the start type" + + # NULL for any elements including char (occurs at least once) + for x in range(3,6): + if not all(item.isdigit() for item in list(retro_list[x])): + retro_list[x] = None + player_instance = { 'player_id': retro_list[1] - , 'player_name': retro_list[2] + , 'player_name': re.sub('"', '', retro_list[2]) , 'player_team': retro_list[3] , 'batting_order': retro_list[4] , 'field_position': retro_list[5] + , 'type': 'start' } return player_instance def parse_play(self, retro_list): "Parse the play type with respect to the play_by_play table" + + # Convert unusual pitch counts to NULL + if any([len(retro_list[4]) != 2, retro_list[4] == '??', any(not item.isdigit() for item in list(retro_list[4]))]): + + retro_list[4] = None + play_instance = { 'inning': retro_list[1] , 'home_team': retro_list[2] @@ -84,57 +108,70 @@ def parse_play(self, retro_list): , 'rts_pitch_count': retro_list[4] , 'rts_pitch_seq': retro_list[5] , 'play_meta': retro_list[6] + , 'type': 'play' } return play_instance def parse_sub(self, retro_list): "Parse the sub type" + sub_instance = { 'player_id': retro_list[1] - , 'player_name': retro_list[2] + , 'player_name': re.sub(']', '', retro_list[2]) , 'player_team': retro_list[3] , 'batting_order': retro_list[4] , 'field_position': retro_list[5] + , 'type': 'sub' } return sub_instance def parse_data(self, retro_list): "Parse the 'data' type" + data_instance = { 'player_id': retro_list[2] , 'earned_runs': retro_list[3] + , 'type': 'data' } return data_instance def parse_ladj(self, retro_list): "Parse the ladj type, also know as BOOT" + ladj_instance = { 'team': retro_list[1] , 'batting_order': retro_list[2] + , 'type': 'ladj' } return ladj_instance def parse_badj(self, retro_list): "Parse the badj type" + badj_instance = { 'player_id': retro_list[1] , 'badj_note': retro_list[2] + , 'type': 'badj' } return badj_instance def parse_padj(self, retro_list): "Parse the padj type" + padj_instance = { 'player_id': retro_list[1] , 'padj_note': retro_list[2] + , 'type': 'padj' } return padj_instance def parse_com(self, retro_list): - return {'comment':retro_list[1]} + "Parse the comment type" + + return {'comment':re.sub('"', '', retro_list[1]), 'type': 'com'} diff --git a/etl/RetrosheetTableBuilder.py b/etl/RetrosheetTableBuilder.py index f8a9c60..31cb797 100644 --- a/etl/RetrosheetTableBuilder.py +++ b/etl/RetrosheetTableBuilder.py @@ -4,122 +4,215 @@ """ import os +import re from os import listdir from os.path import isfile, join import sys import itertools - +import csv +import json import psycopg2 - from etl.RetrosheetParser import RetrosheetLineParser +from etl.RetrosheetTableSetter import RetrosheetTableSetter class RetrosheetTableBuilder(object): - def __init__(self): - conn = psycopg2.connect(**db_credentials) - cur = conn.cursor() - - GAME_INFO_TABLE = 'game_info' - PLAY_BY_PLAY_TABLE = 'play_by_play' - STARTERS_TABLE = 'starters' - - - def build_rts_tables(self, dir_path, tmp_dir): + + # Try to connect to db + try: + self.conn = psycopg2.connect("dbname='patrickfisher' user='patrickfisher' host='localhost'") + print "Connected!" + except: + print "I am unable to connect to the database" + + # Open a cursor to the connection. + self.cur = self.conn.cursor() + + def build_rts_tables(self, dir_path, tmp_dir=''): """From a directory of retrosheet files, build relational tables""" + # Create tables self.create_rts_tables() + # Loop through files for element in listdir(dir_path): - if isfile(join(dir_path, element)) and \ - element.lower().endswith(('.eva', '.evn')): - - self.file_path = join(dir_path, element) - self.rts_line_parser(self.file_path, tmp_dir) - self.rts_db_update(tmp_dir) - self.clean_tmp_dir(tmp_dir) + + if isfile(join(dir_path, element)) and element.lower().endswith(('.eva', '.evn')): + + self.file_path = join(dir_path, element) + print 'Reading ' + self.file_path + rts_output = self.rts_line_parser(self.file_path) + self.rts_db_update(rts_output) + print 'Finished ' + self.file_path +# self.clean_tmp_dir(tmp_dir) else: print('Unrecognized file: %s. Skipping' % (element)) - - + def rts_line_parser(self, file_path, tmp_dir=''): - """From a retrosheet file, parse it into a relational table""" - - rsp = RetrosheetLineParser() - tmp_f = open(join(tmp_dir, 'temp.txt'), 'wb') - tmp_output = csv.writer(tmp_f, quoting = csv.QUOTE_ALL) - with open(self.file_path,'r') as f: - for line in f: - line_parts = line.rstrip().split(',') - parsed_line = rsp.type_parser(line_parts, a.parser_map) - if parsed_line: - parsed_line_list = list(itertools.chain.from_iterable([[k, parsed_line[k]] for k in parsed_line])) - tmp_output.writerow(parsed_line_list) - - def rts_db_update(self, tmp_dir): - """Update the database from a set of temp files in relational format""" - pass - - + """From a retrosheet file, parse it into a relational table""" + + # Create parsing and setting objects + rsp = RetrosheetLineParser() + rts_setter = RetrosheetTableSetter() + + with open(file_path,'r') as f: + for line in f: + # Split up a line + line_parts = line.rstrip().split(',') + + # Parse a line from a retrosheet file + parsed_line = rsp.type_parser(line_parts, rsp.parser_map) + + # Determine how to store a parsed line via rts_setter + if parsed_line: + rts_event_dict = rts_setter.setter_parser(parsed_line, rts_setter.setter_map) + + print "Done with " + file_path + "!" + + return rts_event_dict + + def rts_db_update(self, rts_output): + """Update the database from a set of temp files in relational format""" + + # Add values to starters table + self.cur.executemany(""" + INSERT INTO starters (game_id, player_id, player_name, field_position, batting_order, home_team) + VALUES + (%(game_id)s, %(player_id)s, %(player_name)s, %(field_position)s, %(batting_order)s, %(home_team)s)""", [rts_output['starters'][dict] for dict in rts_output['starters'].keys()]) + self.conn.commit() + + # Dump 'info' types into one json to read to database. + for key in rts_output['ginfo'].keys(): + rts_output['ginfo'][key]['info'] = json.dumps(rts_output['ginfo'][key]['info']) + + # Add values to game_info table + self.cur.executemany(""" + INSERT INTO game_info + (game_id, visteam, hometeam, date, info) + VALUES + (%(game_id)s, %(visteam)s, %(hometeam)s, %(date)s, %(info)s)""", + [rts_output['ginfo'][dict] for dict in rts_output['ginfo'].keys()]) + self.conn.commit() + + # Add values to play_by_play table + self.cur.executemany(""" + INSERT INTO play_by_play + (inning, play_meta, rts_pitch_sequence, rts_play_sequence, rts_pitch_count, player_id, game_id, play_type, bottom_inning, sub_batting_order, sub_field_position, comment) + VALUES + (%(inning)s, %(play_meta)s, %(rts_pitch_sequence)s, %(rts_play_sequence)s, %(rts_pitch_count)s, %(player_id)s, %(game_id)s, %(play_type)s, %(bottom_inning)s, %(sub_batting_order)s, %(sub_field_position)s, %(comment)s)""", + [rts_output['pbp'][dict] for dict in rts_output['pbp'].keys()]) + self.conn.commit() + def clean_tmp_dir(self, tmp_dir): - """Remove files from a tmp directory""" - pass - - + """Remove files from a tmp directory""" + pass + + def check_rts_exist(self, table_name): + """Check whether an rts table exists""" + + self.cur.execute(""" + SELECT EXISTS( + SELECT 1 + FROM information_schema.tables + WHERE table_schema = 'public' + AND table_name = %s); + """, (table_name,)) + ans = self.cur.fetchall() + self.conn.commit() + + return ans[0][0] + + def check_type_exist(self, type): + """Check whether an rts type exists""" + + self.cur.execute(""" + SELECT EXISTS( + SELECT 1 + FROM pg_type + WHERE typname = 'play_type'); + """) + ans = self.cur.fetchall() + self.conn.commit() + + return ans[0][0] + def create_rts_tables(self): - print("Create starters table") - self.cur.execute(""" - CREATE TABLE starters ( - game_id VARCHAR(12) NOT NULL - , player_id VARCHAR(8) NOT NULL - , player_name VARCHAR(35) - , player_team VARCHAR(3) - , home_team BOOLEAN - , batting_order INTEGER - , field_position INTEGER - , PRIMARY KEY (game_id, player_id) - ); - """) - cur.commit() - - print("Create play_by_play table") - self.cur.execute(""" - CREATE TYPE play_type AS ENUM ('play', 'sub', 'badj', 'padj' - , 'ladj') - - CREATE TABLE play_by_play ( - game_id VARCHAR(12) NOT NULL - , player_id VARCHAR(8) NOT NULL - , rts_play_sequence INTEGER NOT NULL - , player_team VARCHAR(3) - , play_type play_type - , inning INTEGER - , bottom_inning BOOLEAN - , home_team BOOLEAN - , rts_pitch_count VARCHAR(3) - , rts_pitch_seq VARCHAR(15) - , play_meta VARCHAR(20) - , sub_batting_order INTEGER - , sub_field_position INTEGER - , PRIMARY KEY (game_id, player_id, rts_play_sequence) - ); - """) - cur.commit() + """Create rts tables""" + + # If a starters table cannot be found, make it. + if self.check_rts_exist('starters') is False: + print('Create starters table') + self.cur.execute(""" + CREATE TABLE starters ( + game_id VARCHAR(12) NOT NULL + , player_id VARCHAR(8) NOT NULL + , player_name VARCHAR(35) + , home_team BOOLEAN + , batting_order INTEGER + , field_position INTEGER + , PRIMARY KEY (game_id, player_id) + ); + """) + self.conn.commit() + else: + + print('starters table already exists!') + + # If play_type type cannot be found, make it. + if self.check_type_exist('play_type') is False: + print("Create play_type TYPE") + self.cur.execute(""" + CREATE TYPE play_type AS ENUM ('play', 'sub', 'badj', 'padj' + , 'ladj'); + """) + self.conn.commit() + else: + + print("play_type type already exists!") + + # If play_by_play table cannot be found, make it. + if self.check_rts_exist('play_by_play') is False: + print("Create play_by_play table") + self.cur.execute(""" + CREATE TABLE play_by_play ( + game_id VARCHAR(12) NOT NULL + , player_id VARCHAR(8) NOT NULL + , rts_play_sequence INTEGER NOT NULL + , play_type play_type + , inning INTEGER + , bottom_inning BOOLEAN + , rts_pitch_count VARCHAR(3) + , rts_pitch_sequence VARCHAR(30) + , play_meta VARCHAR(50) + , sub_batting_order INTEGER + , sub_field_position INTEGER + , comment VARCHAR(500) + , PRIMARY KEY (game_id, player_id, rts_play_sequence) + ); + """) + self.conn.commit() + else: + + print('play_by_play table already exists!') + + # If a game_info table cannot be found, make it. + if self.check_rts_exist('game_info') is False: + print("Create game_info table") + self.cur.execute(""" + CREATE TABLE game_info ( + game_id VARCHAR(12) NOT NULL + , visteam VARCHAR(3) + , hometeam VARCHAR(3) + , date DATE + , info JSON + , PRIMARY KEY (game_id) + ) + """); + self.conn.commit() + else: + print('game_info table already exists!') - print("Create game_info table") - self.cur.execute(""" - - CREATE TABLE game_info ( - game_id VARCHAR(12) NOT NULL - , vistteam VARCHAR(3) - , hometeam VARCHAR(3) - , site VARCHAR(5) - , number INTEGER - , date DATE - , info JSON - , PRIMARY KEY (game_id) - """); - cur.commit() if __name__ == "__main__": """Parse cmd line args and build retrosheet tables""" diff --git a/etl/RetrosheetTableSetter.py b/etl/RetrosheetTableSetter.py new file mode 100644 index 0000000..8148ce5 --- /dev/null +++ b/etl/RetrosheetTableSetter.py @@ -0,0 +1,220 @@ +""" +Retrosheet Table Builder + +Goal: Organize parsed rts data into dict structure that can be read into database +""" + +class RetrosheetTableSetter(object): + + def __init__(self): + + # Read parsed line into appropriate dict variable based on setter_map + self.setter_map = { + "id": self.setter_id, + "version": self.setter_version, + "info": self.setter_info, + "start": self.setter_start, + "play": self.setter_play, + "sub": self.setter_sub, + "data": self.setter_data, + "com": self.setter_com, + "ladj": self.setter_ladj, + "badj": self.setter_badj, + "padj": self.setter_padj + } + + # Dicts to be read into database + self.ginfo = {} + self.pbp = {} + self.starters = {} + + # Relevant tmp variables + self.game_id = '' + self.play_id = '' + self.rts_play_sequence = 0 + self.comment_before = None + self.comment_after = None + + def set_play_id(self, parsed_line=None): + """Update play_id variable""" + + if parsed_line is not None: + self.rts_play_sequence += 1 + + # Create play-specific ID + set_play_id = self.game_id + '_' + str(self.rts_play_sequence) + + return set_play_id + + def setter_parser(self, parsed_line, setter_map): + """Determine which setter should be applied to a parsed line""" + + setter_type = parsed_line['type'] + + # Identify and apply appropriate setter + if setter_map.has_key(setter_type): + setter = self.setter_map[setter_type] + type_setter = setter(parsed_line) + else: + print('Unrecognized key: %s. Skipping' % (setter_type)) + + return type_setter + + def setter_id(self, parsed_line): + """Set id type""" + + # Reading a id line necessarily means it is a new game. + self.rts_play_sequence = 0 + + self.game_id = parsed_line['game_id'] + self.ginfo[self.game_id] = {} + self.ginfo[self.game_id]['game_id'] = self.game_id + self.ginfo[self.game_id]['info'] = {} + + return {'ginfo': self.ginfo, 'pbp': self.pbp, 'starters': self.starters} + + def setter_version(self): + """Set version type""" + + # Indicated on retrosheet.org that the version record is obsolete. + pass + + def setter_info(self, parsed_line): + """Set info type""" + + self.info_special_keys = {'visteam', 'hometeam','date'} + + # Get keys from parsed line. + self.keys = [key for key in parsed_line.keys() if key is not 'type'] + + # Add values from identified keys + for key in self.keys: + self.ginfo[self.game_id]['info'][key] = parsed_line[key] + if key in self.info_special_keys: + self.ginfo[self.game_id][key] = parsed_line[key] + + return {'ginfo': self.ginfo, 'pbp': self.pbp, 'starters': self.starters} + + def setter_start(self, parsed_line): + """Set start type""" + + starter_id = self.game_id + parsed_line['player_id'] + self.starters[starter_id] = {} + self.starters[starter_id]['game_id'] = self.game_id + self.starters[starter_id]['player_id'] = parsed_line['player_id'] + self.starters[starter_id]['player_name'] = parsed_line['player_name'] + self.starters[starter_id]['home_team'] = parsed_line['player_team'] + self.starters[starter_id]['batting_order'] = parsed_line['batting_order'] + self.starters[starter_id]['field_position'] = parsed_line['field_position'] + + return {'ginfo': self.ginfo, 'pbp': self.pbp, 'starters': self.starters} + + def setter_play(self, parsed_line): + """Set play type""" + + # Play-relevant comments almost always come after play is described. This updates comment field where one is detected. + # Some comments are very long. Truncates at 500 characters. + if self.comment_after and len(self.comment_after) > 500: + self.pbp[self.play_id]['comment'] = self.comment_after[:500] + self.comment_after = None + elif self.comment_after and len(self.comment_after) < 500: + self.pbp[self.play_id]['comment'] = self.comment_after + self.comment_after = None + + self.play_id = self.set_play_id(parsed_line) + self.pbp[self.play_id] = {} + self.pbp[self.play_id]['game_id'] = self.game_id + self.pbp[self.play_id]['player_id'] = parsed_line['player_id'] + self.pbp[self.play_id]['rts_play_sequence'] = self.rts_play_sequence + self.pbp[self.play_id]['play_type'] = parsed_line['type'] + self.pbp[self.play_id]['inning'] = parsed_line['inning'] + self.pbp[self.play_id]['bottom_inning'] = parsed_line['home_team'] + self.pbp[self.play_id]['rts_pitch_count'] = parsed_line['rts_pitch_count'] + self.pbp[self.play_id]['rts_pitch_sequence'] = parsed_line['rts_pitch_seq'] + self.pbp[self.play_id]['play_meta'] = parsed_line['play_meta'] + + # Occasionally there is a comment before a play. Handles those instances. + if self.comment_before: + self.pbp[self.play_id]['comment'] = self.comment_before + self.comment_before = None + + self.add_none(self.pbp[self.play_id]) + + return {'ginfo': self.ginfo, 'pbp': self.pbp, 'starters': self.starters} + + def setter_sub(self, parsed_line): + """Set sub type""" + + self.play_id = self.set_play_id(parsed_line) + self.pbp[self.play_id] = {} + self.pbp[self.play_id]['game_id'] = self.game_id + self.pbp[self.play_id]['player_id'] = parsed_line['player_id'] + self.pbp[self.play_id]['rts_play_sequence'] = self.rts_play_sequence + self.pbp[self.play_id]['play_type'] = parsed_line['type'] +# self.pbp[self.play_id]['player_team'] = parsed_line['player_team'] + self.pbp[self.play_id]['sub_batting_order'] = parsed_line['batting_order'] + self.pbp[self.play_id]['sub_field_position'] = parsed_line['field_position'] + self.pbp[self.play_id]['player_name'] = parsed_line['player_name'] + + self.add_none(self.pbp[self.play_id]) + + return {'ginfo': self.ginfo, 'pbp': self.pbp, 'starters': self.starters} + + def setter_data(self, parsed_line): + """Set data type""" + + set_special_keys = ['earned_runs', 'player_id'] + + # Get keys from parsed line. + self.keys = [key for key in parsed_line.keys() if key is not 'type'] + + # Add values from identified keys + if set_special_keys.sort() == self.keys.sort(): + self.ginfo[self.game_id]['info'][parsed_line['player_id']] = parsed_line['earned_runs'] + + return {'ginfo': self.ginfo, 'pbp': self.pbp, 'starters': self.starters} + + def setter_ladj(self, parsed_line): + """Set ladj type""" + + # Inconsistency between how these events are listed in files and described on the website. + pass + + def setter_badj(self, parsed_line): + """Set badj type""" + + self.comment_before = 'badj: ' + parsed_line['badj_note'] + + return {'ginfo': self.ginfo, 'pbp': self.pbp, 'starters': self.starters} + + def setter_padj(self, parsed_line): + + self.comment_before = 'padj: ' + parsed_line['padj_note'] + + return {'ginfo': self.ginfo, 'pbp': self.pbp, 'starters': self.starters} + + def setter_com(self, parsed_line): + """Set com type""" + + # Special case where comment is first event of game. + if self.rts_play_sequence == 0 and self.comment_before is None: + self.comment_before = parsed_line['comment'] + elif self.rts_play_sequence == 0 and self.comment_before: + self.comment_before = self.comment_before + ' ' + parsed_line['comment'] + + # Most comments fit this category. + if self.rts_play_sequence > 0 and self.comment_after is None: + self.comment_after = parsed_line['comment'] + elif self.rts_play_sequence > 0 and self.comment_after: + self.comment_after = self.comment_after + ' ' + parsed_line['comment'] + + return {'ginfo': self.ginfo, 'pbp': self.pbp, 'starters': self.starters} + + def add_none(self, play_id): + """Add None(s) in otherwise unfilled fields""" + + pbp_keys = ['game_id', 'player_id', 'rts_play_sequence', 'play_type', 'inning', 'bottom_inning', 'rts_pitch_count', 'rts_pitch_sequence', 'play_meta', 'sub_batting_order', 'sub_field_position', 'comment'] + + for key in pbp_keys: + if key not in play_id.keys(): + play_id[key] = None diff --git a/tests/test_etl.py b/tests/test_etl.py index 6f16471..4f8cc64 100644 --- a/tests/test_etl.py +++ b/tests/test_etl.py @@ -1,5 +1,6 @@ from etl.RetrosheetParser import RetrosheetLineParser - +import csv +import itertools def test_parse_id(): rsp = RetrosheetLineParser() @@ -105,19 +106,21 @@ def test_parse_com(): def test_type_parser(): - pass -# #Need to specify stable path -# file_path = '/Users/patrickfisher/Documents/mlb/retrosheet/pbp/2013PIT.evn' -# rsp = RetrosheetLineParser() + #Need to specify stable path + file_path = '/Users/patrickfisher/Documents/mlb/retrosheet/pbp/2013PIT.evn' + rsp = RetrosheetLineParser() # tmp_f = open('file.txt', 'wb') # tmp_output = csv.writer(tmp_f, quoting = csv.QUOTE_ALL) -# with open(file_path,'r') as f: -# for line in f: -# line_parts = line.rstrip().split(',') -# parsed_line = rsp.type_parser(line_parts, rsp.parser_map) + with open(file_path,'r') as f: + for line in f: + line_parts = line.rstrip().split(',') + parsed_line = rsp.type_parser(line_parts, rsp.parser_map) + return parsed_line # if parsed_line: -# parsed_line_list = list(itertools.chain.from_iterable([[k, parsed_line[k]] for k in parsed_line])) -# tmp_output.writerow(parsed_line_list) +# # parsed_line_list = list(itertools.chain.from_iterable([[k, parsed_line[k]] for k in parsed_line])) +# parsed_line_list = sum([[k, parsed_line[k]] for k in parsed_line],[]) +# # tmp_output.writerow(parsed_line_list) +# print parsed_line_list test_parse_id() test_parse_info() @@ -129,4 +132,4 @@ def test_type_parser(): test_parse_badj() test_parse_padj() test_parse_com() -test_type_parser() \ No newline at end of file +test_type_parser()