forked from ggtracker/sc2reader
/
sc.py
313 lines (238 loc) · 12.5 KB
/
sc.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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
from __future__ import absolute_import
import os
from sc2reader import utils
from sc2reader import readers
from sc2reader.data import Data_16561,Data_17326,Data_18317,Data_19595
from sc2reader import exceptions
from sc2reader.replay import Replay
from collections import defaultdict
class SC2Reader(object):
"""
The primary interface to the sc2reader library. Acts as a configurable
factory for :class:`Replay` objects. Maintains a set of registered readers,
datapacks, and listeners with filterfuncs that allow the factory to apply
a replay specific context to each replay as it loads.
#TODO: Include some examples here...
See the specific functions below for details.
:param True register_defaults: Automatically registers default readers
and datapacks. Only disable if you know what you are doing.
:param True load_events: Enables parsing of game events. If you are do
not need this information setting to false will reduce replay load
time.
:param True autoplay: Enables auto playing of replays after loading game
events. Playing events triggers enables registered listeners to add
new data features to replays. Option ignored if load_events is false.
:param False load_map: Triggers downloading and parsing of map files
associated with replays as they are loaded. When false, only the map
url and name are available.
:param False verbose: Causes many steps in the replay loading process
to produce more verbose output.
:param string directory: Specifies a base directory to prepend to all paths
before attempting to load the replay.
:param -1 depth: Indicates the maximum search depth when loading replays
from directories. -1 indicates no limit, 0 turns recursion off.
:param list exclude: A list of directory names (not paths) to exclude when
performing recursive searches while loading replays from directories.
:param False followlinks: Enables symlink following when recursing through
directories to load replay files.
"""
default_options = dict(
# General use
verbose=False,
debug=False,
# Related to the SC2Reader class only
register_defaults=True,
# Related to creating new replay objects
autoplay=True,
load_events=True,
# Related to passing paths into load_replay(s)
directory='',
# Related to passing a directory to load_replays
depth=-1,
exclude=[],
followlinks=False
)
def __init__(self, **options):
self.reset()
self.configure(**options)
if self.options.get('register_defaults',None):
self.register_defaults()
def configure(self, **new_options):
"""
Update the factory settings with the specified overrides.
See :class:`SC2Reader` for a list of available options.
:param new_options: Option values to override current factory settings.
"""
self.options.update(new_options)
def reset(self):
"""
Resets the current factory to default settings and removes all
registered readers, datapacks, and listeners.
"""
self.options = utils.AttributeDict(self.default_options)
self.registered_readers = defaultdict(list)
self.registered_datapacks = list()
self.registered_listeners = defaultdict(list)
def load_replays(self, collection, options=None, **new_options):
"""
Loads the specified collection of replays using the current factory
settings with specified overrides.
:param collection: Either a directory path or a mixed collection of
directories, file paths, and open file objects.
:param None options: When options are passed directly into the options
parameter the current factory settings are ignored and only the
specified options are used during replay load.
:param new_options: Options values to override current factory settings
for the collection of replays to be loaded.
:rtype: generator(:class:`Replay`)
"""
options = options or utils.merged_dict(self.options, new_options)
# Get the directory and hide it from nested calls
directory = options.get('directory','')
if 'directory' in options: del options['directory']
if isinstance(collection, basestring):
full_path = os.path.join(directory, collection)
for replay_path in utils.get_replay_files(full_path, **options):
with open(replay_path) as replay_file:
yield self.load_replay(replay_file, options=options)
else:
for replay_file in collection:
if isinstance(replay_file, basestring):
full_path = os.path.join(directory, replay_file)
if os.path.isdir(full_path):
for replay in self.load_replays(full_path, options=options):
yield replay
else:
yield self.load_replay(full_path, options=options)
else:
yield self.load_replay(replay_file, options=options)
def load_replay(self, replay_file, options=None, **new_options):
"""
Loads the specified replay using current factory settings with the
specified overrides.
:param replay: An open file object or a path to a single file.
:param None options: When options are passed directly into the options
parameter the current factory settings are ignored and only the
specified options are used during replay load.
:param new_options: Options values to override current factory settings
while loading this replay.
:rtype: :class:`Replay`
"""
options = options or utils.merged_dict(self.options, new_options)
load_events = options.get('load_events',True)
autoplay = options.get('autoplay',True)
if isinstance(replay_file, basestring):
location = os.path.join(options.get('directory',''), replay_file)
with open(location, 'rb') as replay_file:
return self.load_replay(replay_file, options=options)
if options['verbose']:
print replay_file.name
replay = Replay(replay_file, **options)
for data_file in ('replay.initData',
'replay.details',
'replay.attributes.events',
'replay.message.events',):
reader = self.get_reader(data_file, replay)
replay.read_data(data_file, reader)
replay.load_details()
replay.load_players()
if load_events:
for data_file in ('replay.game.events',):
reader = self.get_reader(data_file, replay)
replay.read_data(data_file, reader)
replay.load_events(self.get_datapack(replay))
if autoplay:
replay.listeners = self.get_listeners(replay)
replay.start()
return replay
def get_reader(self, data_file, replay):
for callback, reader in self.registered_readers[data_file]:
if callback(replay):
return reader
def get_datapack(self, replay):
for callback, datapack in self.registered_datapacks:
if callback(replay):
return datapack
return None
def get_listeners(self, replay):
listeners = defaultdict(list)
for event_class in self.registered_listeners.keys():
for callback, listener in self.registered_listeners[event_class]:
if callback(replay):
listeners[event_class].append(listener)
return listeners
def register_listener(self, events, listener, filterfunc=lambda r: True):
"""
Allows you to specify event listeners for adding new features to the
:class:`Replay` objects on :meth:`~Replay.play`. sc2reader comes with a
small collection of :class:`Listener` classes that you can apply to your
replays as needed.
Events are sent to listeners in registration order as they come up. By
specifying a parent class you can register a listener to a set of events
at once instead of listing them out individually. See the tutorials for
more information.
:param events: A list of event classes you want sent to this listener.
Registration to a single event can be done by specifying a single
event class instead of a list. An isinstance() check is used so
you can catch sets of classes at once by supplying a parent class.
:param listener: The :class:`Listener` object you want events sent to.
:param filterfunc: A function that accepts a partially loaded
:class:`Replay` object as an argument and returns true if the
reader should be used on this replay.
"""
try:
for event in events:
self.registered_listeners[event].append((filterfunc, listener))
except TypeError:
self.registered_listeners[event].append((filterfunc, listener))
def register_reader(self, data_file, reader, filterfunc=lambda r: True):
"""
Allows you to specify your own reader for use when reading the data
files packed into the .SC2Replay archives. Datapacks are checked for
use with the supplied filterfunc in reverse registration order to give
user registered datapacks preference over factory default datapacks.
Don't use this unless you know what you are doing.
:param data_file: The full file name that you would like this reader to
parse.
:param reader: The :class:`Reader` object you wish to use to read the
data file.
:param filterfunc: A function that accepts a partially loaded
:class:`Replay` object as an argument and returns true if the
reader should be used on this replay.
"""
self.registered_readers[data_file].insert(0,(filterfunc, reader))
def register_datapack(self, datapack, filterfunc=lambda r: True):
"""
Allows you to specify your own datapacks for use when loading replays.
Datapacks are checked for use with the supplied filterfunc in reverse
registration order to give user registered datapacks preference over
factory default datapacks.
This is how you would add mappings for your favorite custom map.
:param datapack: A :class:`BaseData` object to use for mapping unit
types and ability codes to their corresponding classes.
:param filterfunc: A function that accepts a partially loaded
:class:`Replay` object as an argument and returns true if the
datapack should be used on this replay.
"""
self.registered_datapacks.insert(0,(filterfunc, datapack))
def register_defaults(self):
"""Registers all factory default objects."""
self.register_default_readers()
self.register_default_datapacks()
def register_default_readers(self):
"""Registers factory default readers."""
self.register_reader('replay.details', readers.DetailsReader_Base())
self.register_reader('replay.initData', readers.InitDataReader_Base())
self.register_reader('replay.message.events', readers.MessageEventsReader_Base())
self.register_reader('replay.attributes.events', readers.AttributesEventsReader_Base(), lambda r: r.build < 17326)
self.register_reader('replay.attributes.events', readers.AttributesEventsReader_17326(), lambda r: r.build >= 17326)
self.register_reader('replay.game.events', readers.GameEventsReader_Base(), lambda r: r.build < 16561)
self.register_reader('replay.game.events', readers.GameEventsReader_16561(), lambda r: 16561 <= r.build < 18574)
self.register_reader('replay.game.events', readers.GameEventsReader_18574(), lambda r: 18574 <= r.build < 19595)
self.register_reader('replay.game.events', readers.GameEventsReader_19595(), lambda r: 19595 <= r.build)
def register_default_datapacks(self):
"""Registers factory default datapacks."""
self.register_datapack(Data_16561, lambda r: 16561 <= r.build < 17326)
self.register_datapack(Data_17326, lambda r: 17326 <= r.build < 18317)
self.register_datapack(Data_18317, lambda r: 18317 <= r.build < 19595)
self.register_datapack(Data_19595, lambda r: 19595 <= r.build)