Fixes parsing of replays taken by spectators and improves 1.2 support

Author: GraylinKim

sc2reader/eventparsers.py

@@ -448,7 +448,11 @@ def load(self,event,bytes):
 class UnknownEventParser_04C6(object):
     def load(self,event,bytes):
         event.name = 'unknown04C6'
-        event.bytes += bytes.skip(16,byteCode=True)
+        event.bytes += bytes.peek(16)
+        block1 = bytes.getBig(4)
+        block2 = bytes.getBig(4)
+        block3 = bytes.getBig(4)
+        block4 = bytes.getBig(4)
         return event
 
 class UnknownEventParser_041C(object):
@@ -457,10 +461,11 @@ def load(self,event,bytes):
         event.bytes += bytes.skip(15,byteCode=True)
         return event
 
-class UnknownEventParser_0418or87(object):
+class UnknownEventParser_0487(object):
     def load(self,event,bytes):
         event.name = 'unknown0418-87'
-        event.bytes += bytes.skip(4,byteCode=True)
+        event.data, databytes = bytes.getBig(4,byteCode=True) #Always 00 00 00 01??
+        event.bytes += databytes
         return event
 
 class UnknownEventParser_0400(object):
@@ -469,9 +474,10 @@ def load(self,event,bytes):
         event.bytes += bytes.skip(10,byteCode=True)
         return event
 
-class UnknownEventParser_043C(object):
+class UnknownEventParser_04XC(object):
     def load(self,event,bytes):
-        event.name = 'unknown043C'
+        event.name = 'unknown04XC'
+        print bytes.peek(20)
         return event
 
 #####################################################

sc2reader/exceptions.py

@@ -0,0 +1,6 @@
+class ParseError(Exception):
+    def __init__(self,message,replay,event,bytes):
+        self.message = message
+        self.replay = replay
+        self.event = event
+        self.bytes = bytes

sc2reader/parsers.py

@@ -4,6 +4,10 @@
 from objects import Attribute,Message,Player,Event
 from eventparsers import *
 from utils import ByteStream
+from exceptions import ParseError
+
+from pprint import PrettyPrinter
+pprint = PrettyPrinter(indent=2).pprint
 
 #################################################
 # Parser Dispatch Functions
@@ -148,9 +152,13 @@ def load(self,replay,filecontents):
                 replay.messages.append(Message(time,playerId,target,text))
 
         recorders = [player for player in replay.players if player and player.recorder==True]
-        if len(recorders) != 1:
+        if len(recorders) >1:
             raise ValueError("There should be 1 and only 1 recorder; %s were found" % len(recorders))
-        replay.recorder = recorders[0]
+        elif len(recorders) == 0:
+            #If there are no recorders, then the recorder must not be a player, spectator or referee then
+            replay.recorder = None
+        else:
+            replay.recorder = recorders[0]
 
 ####################################################
 # replay.game.events parsing classes
@@ -179,7 +187,7 @@ class EventParser(object):
             (UnknownEventParser_04X2(), lambda e: e.code & 0x0F == 2 ),
             (UnknownEventParser_0416(), lambda e: e.code == 0x16 ),
             (UnknownEventParser_04C6(), lambda e: e.code == 0xC6 ),
-            (UnknownEventParser_0418or87(), lambda e: e.code == 0x18 or e.code == 0x87 ),
+            (UnknownEventParser_0487(), lambda e: e.code == 0x87 ),
             (UnknownEventParser_0400(), lambda e: e.code == 0x00 ),],
         0x05: [
             (UnknownEventParser_0589(), lambda e: e.code == 0x89 ),],
@@ -203,20 +211,24 @@ def load(self,replay,filecontents):
             #The following byte completes the unique eventObjectidentifier
             first,eventCode = bytes.getBigInt(1),bytes.getBigInt(1)
             eventType,globalFlag,playerId = first >> 5,first & 0x10,first & 0xF
-            
+
             #Create a barebones event from the gathered information
             event = Event(elapsedTime,eventType,eventCode,
                         globalFlag,playerId,location,eventBytes)
 
-            #Get the parser and load the data into the event
-            replay.events.append(self.getParser(event).load(event,bytes))
+            try:
+                #Get the parser and load the data into the event
+                replay.events.append(self.getParser(event).load(event,bytes))
+            except TypeError as e:
+                raise ParseError(e.message,replay,event,bytes)
 
     def getParser(self,event):
         if event.type not in self.parserMap.keys():
             raise TypeError("Unknown eventType: %s at location %s" % (hex(event.type),event.location))
 
         for parser,accept in self.parserMap[event.type]:
-            if accept(event): return parser
+            if accept(event):
+                return parser
 
         raise TypeError("Unknown event: %s - %s at %s" % (hex(event.type),hex(event.code),event.location))
 
@@ -241,15 +253,25 @@ class EventParser_16561(EventParser):
             (CameraMovementEventParser_18(), lambda e: e.code == 0x18 ),
             (CameraMovementEventParser_X1(), lambda e: e.code & 0x0F == 1 ),],
         0x04: [
-            (UnknownEventParser_04X2(), lambda e: e.code & 0x0F == 2 ),
-            (UnknownEventParser_0416(), lambda e: e.code == 0x16 ),
+            (UnknownEventParser_0487(), lambda e: e.code == 0x87 ),
             (UnknownEventParser_04C6(), lambda e: e.code == 0xC6 ),
-            (UnknownEventParser_0418or87(), lambda e: e.code == 0x18 or e.code == 0x87 ),
-            (UnknownEventParser_0400(), lambda e: e.code == 0x00 ),],
+            (UnknownEventParser_04XC(), lambda e: e.code & 0x0F == 0x0C and e.code >> 4 <= 0x04 ),],
         0x05: [
             (UnknownEventParser_0589(), lambda e: e.code == 0x89 ),],
     }
 
+"""
+            (UnknownEventParser_04X2(), lambda e: e.code & 0x0F == 2 ),
+            (UnknownEventParser_0416(), lambda e: e.code == 0x16 ),
+            (UnknownEventParser_0400(), lambda e: e.code == 0x00 ),
+
+
+
+            (UnknownEventParser_04X2(), lambda e: e.code & 0x0F == 2 ),
+            (UnknownEventParser_0416(), lambda e: e.code == 0x16 ),
+            (UnknownEventParser_0400(), lambda e: e.code == 0x00 ),
+"""
+
 class EventParser_17326(EventParser):
 	parserMap = {
         0x00: [
@@ -271,12 +293,9 @@ class EventParser_17326(EventParser):
             (CameraMovementEventParser_18(), lambda e: e.code == 0x18 ),
             (CameraMovementEventParser_X1(), lambda e: e.code & 0x0F == 1 ),],
         0x04: [
-            (UnknownEventParser_04X2(), lambda e: e.code & 0x0F == 2 ),
-            (UnknownEventParser_0416(), lambda e: e.code == 0x16 ),
+            (UnknownEventParser_0487(), lambda e: e.code == 0x87 ),
             (UnknownEventParser_04C6(), lambda e: e.code == 0xC6 ),
-            (UnknownEventParser_0418or87(), lambda e: e.code == 0x18 or e.code == 0x87 ),
-            (UnknownEventParser_0400(), lambda e: e.code == 0x00 ),
-            (UnknownEventParser_041C(), lambda e: e.code == 0x1C ),],
+            (UnknownEventParser_04XC(), lambda e: e.code & 0x0F == 0x0C and e.code >> 4 <= 0x04 ),],
         0x05: [
             (UnknownEventParser_0589(), lambda e: e.code == 0x89 ),],
     }

sc2reader/replay.py

@@ -9,81 +9,104 @@
 
 class Replay(object):
 
-    def __init__(self,replay):
+    def __init__(self,replay,partialParse=True,fullParse=True):
         #Make sure the file exists and is readable
         if not os.access(replay,os.F_OK):
             raise ValueError("File at '%s' cannot be found" % replay)
         elif not os.access(replay,os.R_OK):
             raise ValueError("File at '%s' cannot be read" % replay)
-            
+        
+        self.file = replay
+        self._parseHeader()
+        self.files = MPQArchive(replay).extract()
+        self.parsed = dict(details=False,attributes=False,messages=False,events=False)
+        
+        #These are quickly parsed files that contain most of the game information
+        #The order is important, I need some way to reinforce it in the future
+        if partialParse or fullparse:
+            self._parseDetails()
+            self._parseAttributes()
+            self._parseMessages()
+        
+        if fullParse:
+            self._parseEvents()
+    
+    def _parseHeader(self):
         #Open up a ByteStream for its contents
-        source = ByteStream(open(replay).read())
+        source = ByteStream(open(self.file).read())
 
         #Check the file type for the MPQ header bytes
         if source.getBig(4) != "4D50511B":
-            raise TypeError("File '%s' is not an MPQ file" % replay)
+            raise TypeError("File '%s' is not an MPQ file" % self.file)
 
         #Extract replay header data
-        self.parseMPQHeader(source)
-        
-        #Parse the rest of the archive
-        self.parseMPQArchive(MPQArchive(replay))
-        
-        #Use the parsed information to reorganize the data
-        self.setDerivedValues()
-    
-    def parseMPQHeader(self,source):
-        #Read the MPQ file header
         max_data_size = source.getLittleInt(4) #possibly data max size
         header_offset = source.getLittleInt(4) #Offset of the second header
         data_size = source.getLittleInt(4)     #possibly data size
 
         #Extract replay attributes from the mpq
         data = source.parseSerializedData()
 
+        #Assign all the relevant information to the replay object
         self.build = data[1][4]
         self.versions = (data[1][1],data[1][2],data[1][3],self.build)
         self.releaseString = "%s.%s.%s.%s" % self.versions
         self.frames,self.seconds = (data[3],data[3]/16)
         self.length = (self.seconds/60,self.seconds%60)
 
-    def parseMPQArchive(self,archive):
-        #Extract the archive files
-        files = archive.extract()
-        eventsFile = files['replay.game.events']
-        detailsFile = files['replay.details']
-        messageFile = files['replay.message.events']
-        attributesFile = files['replay.attributes.events']
-        
-        #Load the details file first to get player information
-        parsers.getDetailParser(self.build).load(self,detailsFile)
-        
-        #Next load the attributes file to fill out players and get team information
-        parsers.getAttributeParser(self.build).load(self,attributesFile)
-        
-        #Finally load the events file to get gameplay data and APM
-        parsers.getEventParser(self.build).load(self,eventsFile)
+    def _parseDetails(self):
+        #Get player and map information
+        parsers.getDetailParser(self.build).load(self,self.files['replay.details'])
+        self.parsed['details'] = True
 
-        #We'll also load up the messages for a peak at what was going on
-        parsers.getMessageParser(self.build).load(self,messageFile)
+    def _parseAttributes(self):
+        #The details file data is required for parsing
+        if not self.parsed['details']:
+            raise ValueError("The replay details must be parsed before parsing attributes")
+            
+        parsers.getAttributeParser(self.build).load(self,self.files['replay.attributes.events'])
+        self.parsed['attributes'] = True
 
-    def setDerivedValues(self):
+        #We can now create teams
         self.teams = defaultdict(list)
         for player in self.players[1:]: #Skip the 'None' player 0
             self.teams[player.team].append(player)
-
+        
+    def _parseMessages(self):
+        #The details file data is required for parsing
+        if not self.parsed['details']:
+            raise ValueError("The replay details must be parsed before parsing messages")
+            
+        parsers.getMessageParser(self.build).load(self,self.files['replay.message.events'])
+        self.parsed['messages'] = True
+        
+    def _parseEvents(self):
+        #The details file data is required for parsing
+        if not self.parsed['details']:
+            raise ValueError("The replay details must be parsed before parsing events")
+            
+        parsers.getEventParser(self.build).load(self,self.files['replay.game.events'])
+        self.parsed['events'] = True
+        
+        #We can now sort events by type and get results
         self.eventsByType = defaultdict(list)
         for event in self.events:
             self.eventsByType[event.name].append(event)
+            
+        self._processResults()
 
-        self.processResults()
-        
-    def processResults(self):
+    def _processResults(self):
+        #The details,attributes, and events are required
+        if not (self.parsed['details'] and self.parsed['attributes'] and self.parsed['events']):
+            raise ValueError("The replay details must be parsed before parsing attributes")
+            
         #Remove players from the teams as they drop out of the game   
         self.results = dict([team,len(players)] for team,players in self.teams.iteritems())
         for event in self.eventsByType['leave']:
-            team = self.players[event.player].team
-            self.results[team] -= 1 
+            #Some spectator actions seem to be recorded, they aren't on teams anyway
+            if event.player < len(self.players):
+                team = self.players[event.player].team
+                self.results[team] -= 1 
 
         #mark all teams with no players left as losing, save the rest of the teams
         remaining = set()
@@ -96,7 +119,7 @@ def processResults(self):
         #If, at the end, only one team remains then that team has won
         if len(remaining) == 1:
             self.results[remaining.pop()] = "Won"
-        else:    
+        elif self.recorder:    
             #The other results are unknown except in the (common) case that the
             #recorder is the last one on his team to leave. In this case, the
             #result for his team can be known