Skip to content

Commit ea503a1

Browse files
GraylinKimGraylinKim
committed
Reorganization of the __init__.py file.
Moves package configuration variables to the config module and moves the read_header utility to the utils module to discourage use (use mpyq instead).
1 parent 67c0ff3 commit ea503a1

File tree

3 files changed

+137
-134
lines changed

3 files changed

+137
-134
lines changed

‎sc2reader/__init__.py‎

Lines changed: 27 additions & 134 deletions
Original file line numberDiff line numberDiff line change
@@ -1,170 +1,63 @@
11
import os
22

3+
from config import FULL, PARTIAL, CUSTOM, FILES, PROCESSORS, READERS
34
from mpyq import MPQArchive
4-
from utils import ReplayBuffer, LITTLE_ENDIAN
55
from objects import Replay
6-
from processors import *
7-
from readers import *
6+
from utils import read_header
87

98
__version__ = "0.3.0"
9+
__author__ = "Graylin Kim <[email protected]>"
1010

11-
FULL = "FULL"
12-
PARTIAL = "PARTIAL"
13-
CUSTOM = "CUSTOM"
14-
15-
FILES = {
16-
"FULL": [
17-
'replay.initData',
18-
'replay.details',
19-
'replay.attributes.events',
20-
'replay.message.events',
21-
'replay.game.events'
22-
],
23-
24-
"PARTIAL": [
25-
'replay.initData',
26-
'replay.details',
27-
'replay.attributes.events',
28-
'replay.message.events'
29-
],
30-
}
31-
32-
PROCESSORS = {
33-
"FULL": [
34-
PeopleProcessor,
35-
AttributeProcessor,
36-
TeamsProcessor,
37-
MessageProcessor,
38-
RecorderProcessor,
39-
EventProcessor,
40-
ApmProcessor,
41-
ResultsProcessor,
42-
],
11+
class SC2Reader(object):
12+
''' Class level interface to sc2reader.
4313
44-
"PARTIAL": [
45-
PeopleProcessor,
46-
AttributeProcessor,
47-
TeamsProcessor,
48-
MessageProcessor,
49-
RecorderProcessor,
50-
],
51-
}
52-
53-
class ReaderMap(object):
54-
def __getitem__(self,key):
55-
if int(key) in (16117,16195,16223,16291):
56-
return {
57-
'replay.initData': InitDataReader(),
58-
'replay.details': DetailsReader(),
59-
'replay.attributes.events': AttributeEventsReader(),
60-
'replay.message.events': MessageEventsReader(),
61-
'replay.game.events': GameEventsReader(),
62-
}
63-
64-
elif int(key) in (16561,16605,16755,16939):
65-
return {
66-
'replay.initData': InitDataReader(),
67-
'replay.details': DetailsReader(),
68-
'replay.attributes.events': AttributeEventsReader(),
69-
'replay.message.events': MessageEventsReader(),
70-
'replay.game.events': GameEventsReader_16561(),
71-
}
72-
73-
elif int(key) in (17326,17682,17811,18092,18221,18317):
74-
return {
75-
'replay.initData': InitDataReader(),
76-
'replay.details': DetailsReader(),
77-
'replay.attributes.events': AttributeEventsReader_17326(),
78-
'replay.message.events': MessageEventsReader(),
79-
'replay.game.events': GameEventsReader_16561(),
80-
}
81-
82-
#This one is also a catch all. If the build isn't recognized, try to use
83-
#the latest parsing code and hope that it works!
84-
elif int(key) in (18574,) or True:
85-
return {
86-
'replay.initData': InitDataReader(),
87-
'replay.details': DetailsReader(),
88-
'replay.attributes.events': AttributeEventsReader_17326(),
89-
'replay.message.events': MessageEventsReader(),
90-
'replay.game.events': GameEventsReader_18574(),
91-
}
92-
93-
READERS = ReaderMap()
94-
95-
def read_header(file):
96-
buffer = ReplayBuffer(file)
97-
98-
#Check the file type for the MPQ header bytes
99-
if buffer.read_hex(4).upper() != "4D50511B":
100-
print "Header Hex was: %s" % buffer.read_hex(4).upper()
101-
raise ValueError("File '%s' is not an MPQ file" % file.name)
102-
103-
#Extract replay header data, we don't actually use this for anything
104-
max_data_size = buffer.read_int(LITTLE_ENDIAN) #possibly data max size
105-
header_offset = buffer.read_int(LITTLE_ENDIAN) #Offset of the second header
106-
data_size = buffer.read_int(LITTLE_ENDIAN) #possibly data size
107-
108-
#Extract replay attributes from the mpq
109-
data = buffer.read_data_struct()
110-
111-
#return the release and frames information
112-
return data[1],data[3]
14+
<<usage documentation here>>
15+
'''
11316

114-
class SC2Reader(object):
11517
def __init__(self, parse="FULL", directory="", processors=[], debug=False, files=None, verbose=False):
116-
#Sanitize the parse level
117-
parse = parse.upper()
118-
if parse not in ("FULL","PARTIAL","CUSTOM"):
18+
try:
19+
#Update and save the reader configuration
20+
parse = parse.upper()
21+
files = FILES.get(parse,files)
22+
processors = PROCESSORS.get(parse,processors)
23+
self.__dict__.update(locals())
24+
except KeyError:
11925
raise ValueError("Unrecognized parse argument `%s`" % parse)
120-
121-
#get our defaults and save preferences
122-
files = FILES.get(parse,files)
123-
processors = PROCESSORS.get(parse,processors)
124-
self.__dict__.update(locals())
125-
26+
12627
def read(self, location):
127-
#Sanitize the location provided (accounting for directory option)
12828
if self.directory:
12929
location = os.path.join(self.directory,location)
130-
if not os.path.exists(location):
131-
raise ValueError("Path `%s` cannot be found" % location)
132-
30+
13331
if self.verbose: print "Reading: %s" % location
134-
135-
#If its a directory, read each subfile/directory and combine the lists
32+
13633
if os.path.isdir(location):
34+
#SC2Reader::read each subfile/directory and combine the lists
13735
read = lambda file: self.read(os.path.join(location,file))
13836
tolist = lambda x: [x] if isinstance(x,Replay) else x
13937
return sum(map(tolist,(read(x) for x in os.listdir(location))),[])
140-
141-
#The primary replay reading routine
38+
14239
else:
143-
if(os.path.splitext(location)[1].lower() != '.sc2replay'):
144-
raise TypeError("Target file must of the SC2Replay file extension")
145-
14640
with open(location) as replay_file:
147-
#Use the MPQ Header information to initialize the replay
14841
release,frames = read_header(replay_file)
14942
replay = Replay(location,release,frames)
15043
archive = MPQArchive(location,listfile=False)
151-
152-
#Extract and Parse the relevant files based on parse level
44+
15345
for file in self.files:
15446
buffer = ReplayBuffer(archive.read_file(file))
15547
READERS[replay.build][file].read(buffer,replay)
156-
157-
#Do cleanup and post processing
48+
15849
for process in self.processors:
15950
replay = process(replay)
160-
51+
16152
return replay
16253

16354
def configure(self,**options):
16455
self.__dict__.update(options)
165-
166-
167-
#Prepare the lightweight interface
56+
57+
58+
''' The package level interface is just a lightweight wrapper around a default
59+
SC2Reader class. See the documentation above for usage details '''
60+
16861
__defaultSC2Reader = SC2Reader()
16962

17063
def configure(**options):

‎sc2reader/config.py‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
from processors import *
2+
from readers import *
3+
4+
FULL = "FULL"
5+
PARTIAL = "PARTIAL"
6+
CUSTOM = "CUSTOM"
7+
8+
FILES = {
9+
"FULL": [
10+
'replay.initData',
11+
'replay.details',
12+
'replay.attributes.events',
13+
'replay.message.events',
14+
'replay.game.events'
15+
],
16+
17+
"PARTIAL": [
18+
'replay.initData',
19+
'replay.details',
20+
'replay.attributes.events',
21+
'replay.message.events'
22+
],
23+
}
24+
25+
PROCESSORS = {
26+
"FULL": [
27+
PeopleProcessor,
28+
AttributeProcessor,
29+
TeamsProcessor,
30+
MessageProcessor,
31+
RecorderProcessor,
32+
EventProcessor,
33+
ApmProcessor,
34+
ResultsProcessor,
35+
],
36+
37+
"PARTIAL": [
38+
PeopleProcessor,
39+
AttributeProcessor,
40+
TeamsProcessor,
41+
MessageProcessor,
42+
RecorderProcessor,
43+
],
44+
}
45+
46+
class ReaderMap(dict):
47+
def __init__(self):
48+
self.set1 = {
49+
'replay.initData': InitDataReader(),
50+
'replay.details': DetailsReader(),
51+
'replay.attributes.events': AttributeEventsReader(),
52+
'replay.message.events': MessageEventsReader(),
53+
'replay.game.events': GameEventsReader(),
54+
}
55+
56+
self.set2 = {
57+
'replay.initData': InitDataReader(),
58+
'replay.details': DetailsReader(),
59+
'replay.attributes.events': AttributeEventsReader(),
60+
'replay.message.events': MessageEventsReader(),
61+
'replay.game.events': GameEventsReader_16561(),
62+
}
63+
64+
self.set3 = {
65+
'replay.initData': InitDataReader(),
66+
'replay.details': DetailsReader(),
67+
'replay.attributes.events': AttributeEventsReader_17326(),
68+
'replay.message.events': MessageEventsReader(),
69+
'replay.game.events': GameEventsReader_16561(),
70+
}
71+
72+
for key in (16117,16195,16223,16291):
73+
self[key] = self.set1
74+
75+
for key in (16561,16605,16755,16939):
76+
self[key] = self.set2
77+
78+
for key in (17326,17682,17811,18092,18221,18317):
79+
self[key] = self.set3
80+
81+
def __getitem__(self,key):
82+
try:
83+
super(self,ReaderMap).__getitem__(key)
84+
except KeyError:
85+
return self.set3
86+
87+
READERS = ReaderMap()

‎sc2reader/utils.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -486,3 +486,26 @@ def key_in_bases(key,bases):
486486
for clazz in set(bases):
487487
if key in clazz.__dict__: return True
488488
return False
489+
490+
def read_header(file):
491+
''' Read the file as a byte stream according to the documentation found at:
492+
http://wiki.devklog.net/index.php?title=The_MoPaQ_Archive_Format
493+
494+
Return the release array and frame count for sc2reader use. For more
495+
defailted header information, access mpyq directly.
496+
'''
497+
buffer = ReplayBuffer(file)
498+
499+
#Sanity check that the input is in fact an MPQ file
500+
if buffer.read_hex(4).upper() != "4D50511B":
501+
print "Header Hex was: %s" % buffer.read_hex(4).upper()
502+
raise ValueError("File '%s' is not an MPQ file" % file.name)
503+
504+
#Extract replay header data, we are unlikely to ever use most of this
505+
max_data_size = buffer.read_int(LITTLE_ENDIAN)
506+
header_offset = buffer.read_int(LITTLE_ENDIAN)
507+
data_size = buffer.read_int(LITTLE_ENDIAN)
508+
header_data = buffer.read_data_struct()
509+
510+
#return the release array (version,major,minor,build) and frame count
511+
return header_data[1],header_data[3]

0 commit comments

Comments
 (0)