Skip to content

Commit 95fcbcc

Browse files
committed
Adds support for file-like objects in sc2reader.read_file.
1 parent af955bf commit 95fcbcc

File tree

1 file changed

+96
-77
lines changed

1 file changed

+96
-77
lines changed

sc2reader/__init__.py

Lines changed: 96 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -87,88 +87,107 @@ def read(self, location, **user_options):
8787
if options.verbose: print "Reading: %s" % location
8888

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

16092
return replays
16193

162-
def read_file(self, file, **options):
163-
replays = self.read(file, **options)
94+
def make_replay(self, replay_file, **options):
95+
options = utils.AttributeDict(options)
16496

165-
# While normal usage would suggest passing in only filenames, it is
166-
# possible that directories could be passed in. Don't fail silently!
167-
if len(replays) > 1:
168-
raise exceptions.MultipleMatchingFilesError(replays)
97+
# The Replay constructor scans the header of the replay file for
98+
# the build number and stores the options for later use. The
99+
# options are copied so subsequent option changes are isolated.
100+
replay_file.seek(0)
101+
replay = objects.Replay(replay_file, **options.copy())
102+
103+
# .SC2Replay files are written in Blizzard's MPQ Archive format.
104+
# The format stores a header which contains a block table that
105+
# specifies the location of each encrypted file.
106+
#
107+
# Unfortunately, some replay sites modify the replay contents to
108+
# add messages promoting their sites without updating the header
109+
# correctly. The listfile option(hack) lets us bypass this issue
110+
# by specifying the files we want instead of generating a list.
111+
#
112+
# In order to wrap mpyq exceptions we have to do this try hack.
113+
try:
114+
replay_file.seek(0)
115+
archive = mpyq.MPQArchive(replay_file, listfile=False)
116+
except KeyboardInterrupt: raise
117+
except:
118+
raise #exceptions.MPQError("Unable to construct the MPQArchive")
119+
120+
# These files are configured for either full or partial parsing
121+
for file in options.files:
122+
123+
# To wrap mpyq exceptions we have to do this try hack.
124+
try:
125+
filedata = archive.read_file(file)
126+
except KeyboardInterrupt: raise
127+
except:
128+
raise exceptions.MPQError("Unable to extract file: {0}".format(file))
129+
130+
# For each file, we build a smart buffer object from the
131+
# utf-8 encoded bitstream that mpyq extracts.
132+
buffer = utils.ReplayBuffer(filedata)
133+
134+
# Each version of Starcraft slightly modifies some portions
135+
# of the format for some files. To work with this, the
136+
# config file has a nested lookup structure of
137+
# [build][file]=>reader which returns the appropriate reader
138+
#
139+
# TODO: Different versions also have different data mappings
140+
# sc2reader doesn't yet handle this difficulty.
141+
#
142+
# Readers use the type agnostic __call__ interface so that
143+
# they can be implemented as functions or classes as needed.
144+
#
145+
# Readers return the extracted information from the buffer
146+
# object which gets stored into the raw data dict for later
147+
# use in post processing because correct interpretation of
148+
# the information often requires data from other files.
149+
reader = config.readers[replay.build][file]
150+
reference_name = '_'.join(file.split('.')[1:])
151+
replay.raw[reference_name] = reader(buffer, replay)
152+
153+
# Now that the replay has been loaded with the "raw" data from
154+
# the archive files we run the system level post processors to
155+
# organize the data into a cross referenced data structure.
156+
#
157+
# After system level processors have run, call each of the post
158+
# processors provided by the user. This would be a good place to
159+
# convert the object to a serialized json string for cross
160+
# language processes or add custom attributes.
161+
#
162+
# TODO: Maybe we should switch this to a hook based architecture
163+
# Needs to be able to load "contrib" type processors..
164+
for process in [processors.Full]+options.processors:
165+
replay = process(replay)
166+
167+
return replay
168+
169+
def read_file(self, file_in, **user_options):
170+
# Support file-like objects (with a read method)
171+
if hasattr(file_in, 'read'):
172+
173+
# Base the options off a copy to leave the Reader options uneffected.
174+
options = self.options.copy()
175+
options.update(user_options)
176+
177+
return self.make_replay(file_in, **options)
178+
179+
# Also support filepath strings
180+
else:
181+
replays = self.read(file_in, **options)
182+
183+
# While normal usage would suggest passing in only filenames, it is
184+
# possible that directories could be passed in. Don't fail silently!
185+
if len(replays) > 1:
186+
raise exceptions.MultipleMatchingFilesError(replays)
187+
188+
# Propogate the replay in a singular context
189+
return replays[0] if len(replays) > 0 else None
169190

170-
# Propogate the replay in a singular context
171-
return replays[0] if len(replays) > 0 else None
172191

173192
"""sc2reader uses a default SC2Reader class instance to provide a package level
174193
interface to its functionality. The package level interface presents the same

0 commit comments

Comments
 (0)