sc2reader now supports all production releases of Starcraft II. Well, parsing passes, lots of missing codes and scramble bytes abound. Also reorganizes the logic for handling the slowly evolving replay format

GraylinKim · GraylinKim · commit 198aa3d3a99a · 2011-05-17T03:54:40.000-04:00
diff --git a/sc2reader/__init__.py b/sc2reader/__init__.py
@@ -1,13 +1,5 @@
 import os
 
-try:
-    from collections import OrderedDict
-except ImportError:
-    from ordereddict import OrderedDict
-except ImportError:
-    from sys import exit
-    exit("OrderedDict required: Upgrade to python2.7 or `pip install ordereddict`")
-    
 from mpyq import MPQArchive
 from utils import ReplayBuffer, LITTLE_ENDIAN
 from objects import Replay
@@ -42,13 +34,48 @@
             MessageProcessor(),
             RecorderProcessor(),
         ]
-READERS = OrderedDict([
-        ('replay.initData', [ReplayInitDataReader()]),
-        ('replay.details', [ReplayDetailsReader()]),
-        ('replay.attributes.events', [AttributeEventsReader_17326(), AttributeEventsReader()]),
-        ('replay.message.events', [MessageEventsReader()]),
-        ('replay.game.events', [GameEventsReader()]),
-    ])
+
+class ReaderMap(object):
+    def __getitem__(self,key):
+        if int(key) in (16117,16195,16223,16291):
+            return {
+                'replay.initData': ReplayInitDataReader(),
+                'replay.details': ReplayDetailsReader(),
+                'replay.attributes.events': AttributeEventsReader(),
+                'replay.message.events': MessageEventsReader(),
+                'replay.game.events': GameEventsReader_16291(),
+            }
+
+        elif int(key) in (16561,16605,16755,16939):
+            return {
+                'replay.initData': ReplayInitDataReader(),
+                'replay.details': ReplayDetailsReader(),
+                'replay.attributes.events': AttributeEventsReader(),
+                'replay.message.events': MessageEventsReader(),
+                'replay.game.events': GameEventsReader(),
+            }
+
+        elif int(key) in (17326,17682,17811,18092,18221,18317):
+            return {
+                'replay.initData': ReplayInitDataReader(),
+                'replay.details': ReplayDetailsReader(),
+                'replay.attributes.events': AttributeEventsReader_17326(),
+                'replay.message.events': MessageEventsReader(),
+                'replay.game.events': GameEventsReader(),
+            }
+
+        elif int(key) in (18574,):
+            return {
+                'replay.initData': ReplayInitDataReader(),
+                'replay.details': ReplayDetailsReader(),
+                'replay.attributes.events': AttributeEventsReader_17326(),
+                'replay.message.events': MessageEventsReader(),
+                'replay.game.events': GameEventsReader_18574(),
+            }
+        else:
+            raise KeyError(key)
+
+READERS = ReaderMap()
     
 def read_header(file):
     buffer = ReplayBuffer(file)
@@ -91,20 +118,15 @@ def __init__(self, parse=FULL, directory="", processors=[], debug=False, files=N
     def read(self, location):
         #account for the directory option
         if self.directory: location = os.path.join(self.directory,location)
-        
+
         if not os.path.exists(location):
             raise ValueError("Location must exist")
         
         #If its a directory, read each subfile/directory and combine the lists
         if os.path.isdir(location):
-            replays = list()
-            for filename in os.listdir(location):
-                replay = self.read(os.path.join(location,filename))
-                if isinstance(replay,list):
-                    replays.extend(replay)
-                else:
-                    replays.append(replay)
-            return replays
+            read = lambda file: self.read(os.path.join(location,file))
+            tolist = lambda x: [x] if isinstance(x,Replay) else x
+            return sum(map(tolist,(read(x) for x in os.listdir(location))),[])
             
         #The primary replay reading routine
         else:
@@ -118,14 +140,9 @@ def read(self, location):
                 archive = MPQArchive(location,listfile=False)
                 
                 #Extract and Parse the relevant files based on parse level
-                for file,readers in READERS.iteritems():
-                    if file in self.files:
-                        for reader in readers:
-                            if reader.reads(replay.build):
-                                reader.read(ReplayBuffer(archive.read_file(file)),replay)
-                                break
-                        else:
-                            raise NotYetImplementedError("No parser was found that accepted the replay file;check configuration")
+                for file in self.files:
+                    reader = READERS[replay.build][file]
+                    reader.read(ReplayBuffer(archive.read_file(file)),replay)
                 
                 #Do cleanup and post processing
                 for processor in self.processors:
diff --git a/sc2reader/parsers.py b/sc2reader/parsers.py
@@ -3,6 +3,7 @@
 from sc2reader.objects import *
 from sc2reader.utils import BIG_ENDIAN,LITTLE_ENDIAN
 from collections import defaultdict
+
 class SetupParser(object):
     def parse_join_event(self, buffer, frames, type, code, pid):
         return PlayerJoinEvent(frames, pid, type, code)
@@ -63,6 +64,7 @@ def parse_ability_event(self, buffer, frames, type, code, pid):
                 return LocationAbilityEvent(frames, pid, type, code, None, location)
             
             elif flag in (0x04,0x07):
+                print "Made it!"
                 h = buffer.read_hex(2)
                 hinge = buffer.read_byte()
                 if hinge & 0x20:
@@ -83,15 +85,17 @@ def parse_ability_event(self, buffer, frames, type, code, pid):
             buffer.skip(10)
             return TargetAbilityEvent(frames, pid, type, code, ability, target)
             
-        elif atype < 0x10: #new to patch 1.3.3
+        elif atype in (0x08,0x0a): #new to patch 1.3.3
             #10 bytes total, coordinates have a different format?
             #X coordinate definitely is the first byte, with (hopefully) y next
+            print hex(flag)
             event = UnknownAbilityEvent(frames, pid, type, code, None)
             event.location1 = buffer.read_coordinate()
             buffer.skip(5)
             return event
-        
+
         else:
+            print hex(atype)
             print hex(buffer.cursor)
             raise TypeError()
         
@@ -161,7 +165,108 @@ def read_resource(buffer):
         buffer.skip(8)
         
         return ResourceTransferEvent(frames, pid, type, code, target, minerals, vespene)
+
+class ActionParser_18574(ActionParser):
+    def parse_ability_event(self, buffer, frames, type, code, pid):
+        flag = buffer.read_byte()
+        atype = buffer.read_byte()
         
+        if atype & 0x20: # command card
+            end = buffer.peek(35)
+            ability = buffer.read_byte() << 8 | buffer.read_byte()
+
+            if flag in (0x29, 0x19, 0x14): # cancels
+                # creation autoid number / object id
+                ability = ability << 8 | buffer.read_byte()
+                created_id = buffer.read_object_id()
+                # TODO : expose the id
+                return AbilityEvent(frames, pid, type, code, ability)
+
+            else:
+                ability_flags = buffer.shift(6)
+                ability = ability << 8 | ability_flags
+
+                if ability_flags & 0x10:
+                    # ability(3), coordinates (4), ?? (4)
+                    location = buffer.read_coordinate()
+                    buffer.skip(4)
+                    return LocationAbilityEvent(frames, pid, type, code, ability, location)
+
+                elif ability_flags & 0x20:
+                    # ability(3), object id (4),  object type (2), ?? (10)
+                    code = buffer.read_short() # code??
+                    obj_id = buffer.read_object_id()
+                    obj_type = buffer.read_object_type()
+                    target = (obj_id, obj_type,)
+                    switch = buffer.read_byte()
+                    buffer.read_hex(9)
+                    return TargetAbilityEvent(frames, pid, type, code, ability, target)
+
+                else:
+                    return AbilityEvent(frames,pid,type,code,None)
+
+        elif atype & 0x40: # location/move ??
+            h = buffer.read_hex(2)
+            hinge = buffer.read_byte()
+            if hinge & 0x20:
+                "\t%s - %s" % (hex(hinge),buffer.read_hex(9))
+            elif hinge & 0x40:
+                "\t%s - %s" % (hex(hinge),buffer.read_hex(18))
+            elif hinge < 0x10:
+                pass
+
+            return UnknownLocationAbilityEvent(frames, pid, type, code, None)
+
+        elif atype & 0x80: # right-click on target?
+            # ability (2), object id (4), object type (2), ?? (10)
+            ability = buffer.read_byte() << 8 | buffer.read_byte()
+            obj_id = buffer.read_object_id()
+            obj_type = buffer.read_object_type()
+            target = (obj_id, obj_type,)
+            buffer.skip(10)
+            return TargetAbilityEvent(frames, pid, type, code, ability, target)
+
+        elif atype < 0x10: #new to patch 1.3.3, location now??
+            #10 bytes total, coordinates have a different format?
+            #X coordinate definitely is the first byte, with (hopefully) y next
+            location = buffer.read_coordinate()
+            buffer.skip(5)
+            return LocationAbilityEvent(frames, pid, type, code, None, location)
+
+        else:
+            print hex(atype)
+            print hex(buffer.cursor)
+            raise TypeError()
+
+        print "%s - %s" % (hex(atype),hex(flag))
+        raise TypeError("Shouldn't be here EVER!")
+
+class ActionParser_16291(ActionParser):
+    def parse_ability_event(self, buffer, frames, type, code, pid):
+        buffer.skip(7)
+        switch = buffer.read_byte()
+        if switch in (0x30,0x50):
+            buffer.read_byte()
+        buffer.skip(24)
+        return AbilityEvent(frames, pid, type, code, None)
+
+    def parse_selection_event(self, buffer, frames, type, code, pid):
+        bank = code >> 4
+        selFlags = buffer.read_byte()
+        dsuCount = buffer.read_byte()
+        buffer.read(bits=dsuCount)
+
+        # <count> (<type_id>, <count>,)*
+        object_types = [ (buffer.read_object_type(read_modifier=True), buffer.read_byte(), ) for i in range(buffer.read_byte()) ]
+        # <count> (<object_id>,)*
+        object_ids = [ buffer.read_object_id() for i in range(buffer.read_byte()) ]
+
+        # repeat types count times
+        object_types = chain(*[[object_type,]*count for (object_type, count) in object_types])
+        objects = zip(object_ids, object_types)
+
+        return AbilityEvent(frames, pid, type, code, None)
+
 class Unknown2Parser(object):
     def parse_0206_event(self, buffer, frames, type, code, pid):
         buffer.skip(8)
diff --git a/sc2reader/readers.py b/sc2reader/readers.py
@@ -3,36 +3,11 @@
 from sc2reader.parsers import *
 from sc2reader.objects import *
 from sc2reader.utils import LITTLE_ENDIAN, BIG_ENDIAN
-from sc2reader.utils import key_in_bases, timestamp_from_windows_time
+from sc2reader.utils import timestamp_from_windows_time
 
-#####################################################
-# Metaclass used to help enforce the usage contract
-#####################################################
-class MetaReader(type):
-    def __new__(meta, class_name, bases, class_dict):
-        if class_name != "Reader": #Parent class is exempt from checks
-            assert 'file' in class_dict or key_in_bases('file',bases), \
-                "%s must define the name of the file it reads" % class_name
-
-            assert 'reads' in class_dict or key_in_bases('reads',bases), \
-                "%s must define the 'boolean reads(self, build)' member" % class_name
-
-            assert 'read' in class_dict or key_in_bases('read',bases), \
-                "%s must define the 'void read(self, buffer, replay)' member" % class_name
-
-        return type.__new__(meta, class_name, bases, class_dict)
-
-class Reader(object):
-    __metaclass__ = MetaReader
-		
 #################################################
 
-class ReplayInitDataReader(Reader):
-    file = 'replay.initData'
-
-    def reads(self, build):
-        return True
-        
+class ReplayInitDataReader(object):
     def read(self, buffer, replay):
         
         # Game clients
@@ -57,11 +32,7 @@ def read(self, buffer, replay):
             
 #################################################
 
-class AttributeEventsReader(Reader):
-    file = 'replay.attributes.events'
-    def reads(self, build):
-        return build < 17326
-        
+class AttributeEventsReader(object):
     def read(self, buffer, replay):
         self.load_header(replay, buffer)
         
@@ -78,20 +49,12 @@ def load_header(self, replay, buffer):
         buffer.read_chars(4)
 
 class AttributeEventsReader_17326(AttributeEventsReader):
-    def reads(self, build):
-        return build >= 17326
-
     def load_header(self, replay, buffer):
         buffer.read_chars(5)
         
 ##################################################
 
-class ReplayDetailsReader(Reader):
-    file = 'replay.details'
-
-    def reads(self, build):
-        return True
-    
+class ReplayDetailsReader(object):
     def read(self, buffer, replay):
         data = buffer.read_data_struct()
 
@@ -143,12 +106,7 @@ def read(self, buffer, replay):
 
 ##################################################
 
-class MessageEventsReader(Reader):
-    file = 'replay.message.events'
-
-    def reads(self, build):
-        return True
-    
+class MessageEventsReader(object):
     def read(self, buffer, replay):
         replay.messages, time = list(), 0
 
@@ -183,10 +141,7 @@ def read(self, buffer, replay):
 
 ####################################################
 
-class GameEventsBase(Reader):
-    file = 'replay.game.events'
-    def reads(self, build): return False
-    
+class GameEventsBase(object):
     def read(self, buffer, replay):
         replay.events, frames = list(), 0
         
@@ -206,7 +161,7 @@ def read(self, buffer, replay):
             #print frames
             pid = buffer.shift(5)
             type, code = buffer.shift(3), buffer.read_byte()
-            
+            #print "%s - %s" % (hex(type),hex(code))
             parser = PARSERS[type](code)
             
             if parser == None:
@@ -250,6 +205,10 @@ def get_unknown4_parser(self, code):
         elif code & 0x0F == 0x0C: return self.parse_04XC_event
         
 class GameEventsReader(GameEventsBase,Unknown2Parser,Unknown4Parser,ActionParser,SetupParser,CameraParser):
+    pass
+
+class GameEventsReader_16291(GameEventsBase,Unknown2Parser,Unknown4Parser,ActionParser_16291,SetupParser,CameraParser):
+    pass
 
-    def reads(self, build):
-        return True
+class GameEventsReader_18574(GameEventsBase,Unknown2Parser,Unknown4Parser,ActionParser_18574,SetupParser,CameraParser):
+    pass
