-
Notifications
You must be signed in to change notification settings - Fork 9
/
output.py
123 lines (111 loc) · 4.59 KB
/
output.py
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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
import json
import os
from csv import DictWriter
from typing import Dict
from data import DataObject
from formulas import StatRelevantData
class SaveCsv:
def __init__(self, run_name: str, category_name: str, last_saved_update: Dict[str, str]):
self.object_dir = f"object_data/{run_name}"
os.makedirs(self.object_dir, exist_ok=True)
self.final_filename = f"roll_data/{run_name}-{category_name}.csv"
self.partial_filename = f"{self.final_filename}.partial"
# Created when first row is written
self.file = None
self.csv = None
self.last_saved_object = last_saved_update
def write(
self,
event_type: str,
roll: float,
passed: bool,
update,
is_strike: bool,
strike_roll: float,
strike_threshold: float,
fielder_roll,
baserunners_next,
meta: StatRelevantData,
save_objects: Dict[str, DataObject],
event_time,
):
# fmt: off
row = {
"event_type": event_type,
"event_time": event_time,
"roll": roll,
"passed": passed,
"batting_team_hype": save_objects['stadium'].hype if not update["topOfInning"] else 0,
"pitching_team_hype": save_objects['stadium'].hype if update["topOfInning"] else 0,
"game_id": update["id"],
"play_count": update["playCount"],
"ball_count": update["atBatBalls"],
"strike_count": update["atBatStrikes"],
"out_count": update["halfInningOuts"],
"home_score": update["homeScore"],
"away_score": update["awayScore"],
"inning": update["inning"],
"baserunner_count": update["baserunnerCount"],
"baserunners": str(update["basesOccupied"]),
"baserunners_next": str(baserunners_next),
"is_strike": is_strike,
"strike_roll": strike_roll,
"strike_threshold": strike_threshold,
"fielder_roll": fielder_roll,
"batter_consecutive_hits": save_objects['batter'].consecutive_hits,
"weather": meta.weather,
"season": meta.season,
"day": meta.day,
"runner_count": meta.runner_count,
"top_of_inning": meta.top_of_inning,
"is_maximum_blaseball": meta.is_maximum_blaseball,
"batter_at_bats": meta.batter_at_bats,
}
# fmt: on
for save_key, obj in save_objects.items():
# S20 (1-based) spends a ridiculous amount of time repeatedly
# converting NullPlayer to JSON without this check.
# I think it's because of KLoNGs whose stats aren't available yet.
if obj.id == None and obj.last_update_time == "1970-01-01T00:00:00.000Z":
continue
file_path = f"{self.object_dir}/{obj.id}-{obj.last_update_time}.json".replace(":", "_")
if obj.id not in self.last_saved_object or self.last_saved_object[obj.id] != obj.last_update_time:
to_save = {
"type": obj.object_type,
"last_update_time": obj.last_update_time,
# yes, there's JSON in the JSON. yo dawg
"data": obj.to_json(),
}
with open(file_path, "w") as f:
json.dump(to_save, f)
self.last_saved_object[obj.id] = obj.last_update_time
row[save_key + "_file"] = file_path
if self.csv is None:
self.file = open(self.partial_filename, "w", newline="", encoding="utf-8")
self.csv = DictWriter(self.file, fieldnames=list(row.keys()), extrasaction="ignore")
self.csv.writeheader()
self.csv.writerow(row)
def writeItem(
self, element: str, roll: float, season: int, day: int, roll_type: str, category: str, prefix_index: int = -1
):
row = {
"season": season,
"day": day,
"roll_type": roll_type,
"category": category,
"element": element,
"roll": roll,
"prefix_index": prefix_index,
}
if self.csv is None:
self.file = open(self.partial_filename, "w", newline="", encoding="utf-8")
self.csv = DictWriter(self.file, fieldnames=list(row.keys()), extrasaction="ignore")
self.csv.writeheader()
self.csv.writerow(row)
def close(self):
if not self.file:
return
self.file.close()
self.file = None
self.csv = None
os.replace(self.partial_filename, self.final_filename)