Skip to content

Commit b3dc903

Browse files
GraylinKimGraylinKim
committed
Major modifications as part of a internal redesign.
1 parent b714397 commit b3dc903

File tree

8 files changed

+1064
-1117
lines changed

8 files changed

+1064
-1117
lines changed

‎sc2reader/__init__.py‎

Lines changed: 145 additions & 194 deletions
Original file line numberDiff line numberDiff line change
@@ -1,217 +1,168 @@
1-
# -*- coding: utf-8 -*-
2-
"""
3-
sc2reader
4-
----------
1+
from __future__ import absolute_import
52

6-
A Starcraft II replay parsing library intended to promote innovation in
7-
Starcraft tools and communities. Eventually, it will cover all official
8-
releases and dump easily to JSON for inter-language portability.
3+
import os
4+
from sc2reader import utils
5+
from sc2reader import readers
6+
from sc2reader import data
7+
from sc2reader import exceptions
8+
from sc2reader.replay import Replay
9+
from collections import defaultdict
10+
11+
class SC2Reader(object):
12+
13+
default_options = dict(
14+
# General use
15+
verbose=False,
16+
debug=False,
17+
18+
# Related to the SC2Reader class only
19+
register_defaults=True,
20+
21+
# Related to creating new replay objects
22+
autoplay=True,
23+
complete=True,
24+
25+
# Related to passing paths into load_replay(s)
26+
directory='',
27+
28+
# Related to passing a directory to load_replays
29+
depth=-1,
30+
exclude=[],
31+
followlinks=False
32+
)
33+
34+
def __init__(self, **options):
35+
self.registered_readers = defaultdict(list)
36+
self.registered_datapacks = list()
37+
self.registered_listeners = defaultdict(list)
38+
39+
self.options = utils.AttributeDict(utils.merged_dict(self.default_options, options))
40+
41+
if self.options.register_defaults:
42+
self.register_defaults()
43+
44+
def load_replays(self, replay_collection, options=None, **new_options):
45+
options = options or utils.merged_dict(self.options, new_options)
46+
if isinstance(replay_collection, basestring):
47+
directory = os.path.join(options.get('directory',''), replay_collection)
48+
del options['directory'] # don't need this anymore on this request
49+
for replay_path in utils.get_replay_files(directory, **options):
50+
with open(replay_path) as replay_file:
51+
try:
52+
yield self.load_replay(replay_file, options=options)
53+
except exceptions.MPQError as e:
54+
print e
955

10-
sc2reader has been carefully written for clarity to serve as a starting
11-
point for those who want to write their own parsers, potentially in other
12-
languages for native access.
56+
else:
57+
for replay_file in replay_collection:
58+
yield self.load_replay(replay_file, options=options)
1359

14-
Enjoy.
60+
def load_replay(self, replay_file, options=None, **new_options):
61+
options = options or utils.merged_dict(self.options, new_options)
62+
complete = options.get('complete',True)
63+
autoplay = options.get('autoplay',True)
1564

16-
:copyright: (c) 2011 Graylin Kim
17-
:license: MIT, See LICENSE.txt for details
18-
"""
19-
from __future__ import absolute_import
2065

21-
__version__ = '0.3.3'
66+
if isinstance(replay_file, basestring):
67+
location = os.path.join(options.get('directory',''), replay_file)
68+
with open(location, 'rb') as replay_file:
69+
return self.load_replay(replay_file, options=options)
2270

71+
if options['verbose']:
72+
print replay_file.name
2373

24-
#System imports
25-
import os
74+
replay = Replay(replay_file, **options)
2675

27-
#PyPi imports
28-
import mpyq
29-
30-
#Package imports
31-
from sc2reader import config, objects, utils, processors, exceptions
32-
33-
34-
class Reader(object):
35-
""" The SC2Reader class acts as a factory class for replay objects. The
36-
class accepts a key=value list of options to override defaults (see
37-
config.py) and exposes a very simple read/configure interface and
38-
orchestrates the replay build process.
39-
"""
40-
41-
def __init__(self, **user_options):
42-
""" The constructor makes a copy of the default_options to make sure the
43-
option configuration doesn't propagate back to the default_options.
44-
It should support any arbitrary number of different Reader objects.
45-
"""
46-
self.options = config.default_options.copy()
47-
self.configure(**user_options)
48-
49-
def configure(self, **options):
50-
self.options.update(options)
51-
52-
def read(self, location, **user_options):
53-
""" Read indicated file or recursively read matching files from the
54-
specified directory. Returns a replay or a list of replays depending
55-
on the context.
56-
"""
57-
58-
# Base the options off a copy to leave the Reader options uneffected.
59-
options = self.options.copy()
60-
options.update(user_options)
61-
62-
# The directory option allows users to specify file locations relative
63-
# to a location other than the present working directory by joining the
64-
# location with the directory of their choice.
65-
if options.directory:
66-
location = os.path.join(options.directory, location)
67-
68-
# When passed a directory as the location, the Reader recursively builds
69-
# a list of replays to return using the utils.get_files function. This
70-
# function respects the following arguments:
71-
# * depth: The maximum depth to traverse. Defaults to unlimited (-1)
72-
# * follow_symlinks: Boolean for following symlinks. Defaults to True
73-
# * exclude_dirs: A list of directory names to skip while recursing
74-
# * incldue_regex: A regular expression rule which all returned file
75-
# names must match. Defaults to None
76-
#
77-
replays, files = list(), utils.get_files(location, **options)
78-
79-
# If no files are found, it could be for a variety of reasons
80-
# raise a NoMatchingFilesError to alert them to the situation
81-
if not files:
82-
raise exceptions.NoMatchingFilesError()
83-
84-
for location in files:
85-
if options.verbose: print "Reading: %s" % location
76+
for data_file in ('replay.initData',
77+
'replay.details',
78+
'replay.attributes.events',
79+
'replay.message.events',):
80+
reader = self.get_reader(data_file, replay)
81+
replay.read_data(data_file, reader)
8682

87-
with open(location, 'rb') as replay_file:
88-
replays.append(self.make_replay(replay_file, **options))
89-
90-
return replays
91-
92-
def make_replay(self, replay_file, **options):
93-
options = utils.AttributeDict(options)
94-
95-
# The Replay constructor scans the header of the replay file for
96-
# the build number and stores the options for later use. The
97-
# options are copied so subsequent option changes are isolated.
98-
replay_file.seek(0)
99-
replay = objects.Replay(replay_file, **options.copy())
100-
101-
# .SC2Replay files are written in Blizzard's MPQ Archive format.
102-
# The format stores a header which contains a block table that
103-
# specifies the location of each encrypted file.
104-
#
105-
# Unfortunately, some replay sites modify the replay contents to
106-
# add messages promoting their sites without updating the header
107-
# correctly. The listfile option(hack) lets us bypass this issue
108-
# by specifying the files we want instead of generating a list.
109-
#
110-
# In order to wrap mpyq exceptions we have to do this try hack.
111-
try:
112-
replay_file.seek(0)
113-
archive = mpyq.MPQArchive(replay_file, listfile=False)
114-
except KeyboardInterrupt: raise
115-
except:
116-
raise exceptions.MPQError("Unable to construct the MPQArchive")
117-
118-
# These files are configured for either full or partial parsing
119-
for file in options.files:
120-
121-
# To wrap mpyq exceptions we have to do this try hack.
122-
try:
123-
# Handle the tampering with the message.events file that some sites do
124-
if file == 'replay.message.events':
125-
try:
126-
filedata = archive.read_file(file, force_decompress=True)
127-
except IndexError as e:
128-
# Catch decompression errors
129-
if str(e) == "string index out of range":
130-
filedata = archive.read_file(file, force_decompress=False)
131-
else:
132-
raise
133-
else:
134-
filedata = archive.read_file(file)
135-
except KeyboardInterrupt: raise
136-
except:
137-
raise exceptions.MPQError("Unable to extract file: {0}".format(file))
138-
139-
# For each file, we build a smart buffer object from the
140-
# utf-8 encoded bitstream that mpyq extracts.
141-
buffer = utils.ReplayBuffer(filedata)
142-
143-
# Each version of Starcraft slightly modifies some portions
144-
# of the format for some files. To work with this, the
145-
# config file has a nested lookup structure of
146-
# [build][file]=>reader which returns the appropriate reader
147-
#
148-
# TODO: Different versions also have different data mappings
149-
# sc2reader doesn't yet handle this difficulty.
150-
#
151-
# Readers use the type agnostic __call__ interface so that
152-
# they can be implemented as functions or classes as needed.
153-
#
154-
# Readers return the extracted information from the buffer
155-
# object which gets stored into the raw data dict for later
156-
# use in post processing because correct interpretation of
157-
# the information often requires data from other files.
158-
reader = config.readers[replay.build][file]
159-
reference_name = '_'.join(file.split('.')[1:])
160-
replay.raw[reference_name] = reader(buffer, replay)
161-
162-
# Now that the replay has been loaded with the "raw" data from
163-
# the archive files we run the system level post processors to
164-
# organize the data into a cross referenced data structure.
165-
#
166-
# After system level processors have run, call each of the post
167-
# processors provided by the user. This would be a good place to
168-
# convert the object to a serialized json string for cross
169-
# language processes or add custom attributes.
170-
#
171-
# TODO: Maybe we should switch this to a hook based architecture
172-
# Needs to be able to load "contrib" type processors..
173-
for process in [processors.Full]+options.processors:
174-
replay = process(replay)
83+
replay.load_details()
84+
replay.load_players()
85+
86+
if complete:
87+
for data_file in ('replay.game.events',):
88+
reader = self.get_reader(data_file, replay)
89+
replay.read_data(data_file, reader)
90+
91+
replay.load_events(self.get_datapack(replay))
92+
93+
if autoplay:
94+
replay.listeners = self.get_listeners(replay)
95+
replay.start()
17596

17697
return replay
17798

178-
def read_file(self, file_in, **user_options):
179-
# Support file-like objects (with a read method)
180-
if hasattr(file_in, 'read') and hasattr(file_in, 'seek'):
18199

182-
# Base the options off a copy to leave the Reader options unaffected.
183-
options = self.options.copy()
184-
options.update(user_options)
100+
def get_reader(self, data_file, replay):
101+
for callback, reader in self.registered_readers[data_file]:
102+
if callback(replay):
103+
return reader
185104

186-
return self.make_replay(file_in, **options)
105+
def get_datapack(self, replay):
106+
for callback, datapack in self.registered_datapacks:
107+
if callback(replay):
108+
return datapack
109+
return None
187110

188-
# Also support filepath strings
189-
else:
190-
replays = self.read(file_in, **user_options)
111+
def get_listeners(self, replay):
112+
listeners = defaultdict(list)
113+
for event_class in self.registered_listeners.keys():
114+
for callback, listener in self.registered_listeners[event_class]:
115+
if callback(replay):
116+
listeners[event_class].append(listener)
117+
return listeners
118+
119+
120+
def register_listener(self, event_class, listener, callback=lambda r: True):
121+
self.registered_listeners[event_class].append((callback, listener))
122+
123+
def register_reader(self, data_file, reader, callback=lambda r: True):
124+
self.registered_readers[data_file].insert(0,(callback, reader))
125+
126+
def register_datapack(self, datapack, callback=lambda r: True):
127+
self.registered_datapacks.insert(0,(callback, datapack))
128+
129+
130+
def register_defaults(self):
131+
self.register_default_readers()
132+
self.register_default_datapacks()
133+
self.register_default_listeners()
191134

192-
# While normal usage would suggest passing in only filenames, it is
193-
# possible that directories could be passed in. Don't fail silently!
194-
if len(replays) > 1:
195-
raise exceptions.MultipleMatchingFilesError(replays)
135+
def register_default_readers(self):
136+
self.register_reader('replay.details', readers.DetailsReader_Base())
137+
self.register_reader('replay.initData', readers.InitDataReader_Base())
138+
self.register_reader('replay.message.events', readers.MessageEventsReader_Base())
139+
self.register_reader('replay.attributes.events', readers.AttributesEventsReader_Base(), lambda r: r.build < 17326)
140+
self.register_reader('replay.attributes.events', readers.AttributesEventsReader_17326(), lambda r: r.build >= 17326)
141+
self.register_reader('replay.game.events', readers.GameEventsReader_Base(), lambda r: r.build < 16561)
142+
self.register_reader('replay.game.events', readers.GameEventsReader_16561(), lambda r: 16561 <= r.build < 18574)
143+
self.register_reader('replay.game.events', readers.GameEventsReader_18574(), lambda r: 18574 <= r.build < 19595)
144+
self.register_reader('replay.game.events', readers.GameEventsReader_19595(), lambda r: 19595 <= r.build)
196145

197-
# Propogate the replay in a singular context
198-
return replays[0] if len(replays) > 0 else None
146+
def register_default_datapacks(self):
147+
self.register_datapack(data.Data_16561, lambda r: 16561 <= r.build < 17326)
148+
self.register_datapack(data.Data_17326, lambda r: 17326 <= r.build < 18317)
149+
self.register_datapack(data.Data_18317, lambda r: 18317 <= r.build < 19595)
150+
self.register_datapack(data.Data_19595, lambda r: 19595 <= r.build)
199151

152+
def register_default_listeners(self):
153+
#self.register_listener(Event, PrintEventListener())
154+
pass
200155

201-
"""sc2reader uses a default SC2Reader class instance to provide a package level
202-
interface to its functionality. The package level interface presents the same
203-
functional interface, it just saves the hassle of creating the class object.
204-
"""
205-
__defaultReader = Reader()
206156

207-
def read_file(location, **user_options):
208-
return __defaultReader.read_file(location, **user_options)
157+
__defaultSC2Reader = SC2Reader()
209158

210-
def read(location, **user_options):
211-
return __defaultReader.read(location, **user_options)
159+
register_datapack = __defaultSC2Reader.register_datapack
160+
register_listener = __defaultSC2Reader.register_listener
161+
register_reader = __defaultSC2Reader.register_reader
212162

213-
def configure(**options):
214-
config.default_options.update(options)
163+
get_listeners = __defaultSC2Reader.get_listeners
164+
get_datapack = __defaultSC2Reader.get_datapack
165+
get_reader = __defaultSC2Reader.get_reader
215166

216-
def reset():
217-
__defaultReader = Reader()
167+
load_replay = __defaultSC2Reader.load_replay
168+
load_replays = __defaultSC2Reader.load_replays

0 commit comments

Comments
 (0)