Merge pull request #101 from AoEInsights/fix/scenario-parsing

fixed parsing of scenario matches

mgz/fast/header.py

@@ -34,6 +34,12 @@ def aoc_string(data):
     return data.read(length)
 
 
+def int_prefixed_string(data):
+    """Read length prefixed (4 byte) string."""
+    length = unpack('<I', data)
+    return data.read(length)
+
+
 def de_string(data):
     """Read DE string."""
     assert data.read(2) == b'\x60\x0a'
@@ -229,8 +235,11 @@ def parse_map(data, version):
 def parse_scenario(data, num_players, version, save):
     """Parse scenario section."""
     data.read(4455)
+    scenario_filename = None
     if version is Version.DE:
-        data.read(128)
+        data.read(102)
+        scenario_filename = aoc_string(data)
+        data.read(24)
     instructions = aoc_string(data)
     for _ in range(0, 9):
         aoc_string(data)
@@ -270,15 +279,45 @@ def parse_scenario(data, num_players, version, save):
             settings_version = 2.4
         else:
             settings_version = 2.2
-        end = remainder.find(struct.pack('<d', settings_version))
-        end += 1045
+        end = remainder.find(struct.pack('<d', settings_version)) + 8
     else:
         end = remainder.find(b'\x9a\x99\x99\x99\x99\x99\xf9\x3f') + 13
     data.seek(end - len(remainder), 1)
+
+    if version is Version.DE:
+        data.read(1)
+        n_triggers = unpack("<I", data)
+
+        for _ in range(n_triggers):
+            data.read(22)
+            data.read(4)
+
+            description = int_prefixed_string(data)
+            name = int_prefixed_string(data)
+            short_description = int_prefixed_string(data)
+
+            n_effects = unpack("<I", data)
+
+            for _ in range(n_effects):
+                data.read(216)
+
+                text = int_prefixed_string(data)
+                sound = int_prefixed_string(data)
+
+            data.read(n_effects * 4)
+            n_condition = unpack("<I", data)
+
+            data.read(n_condition * 125)
+
+        trigger_list_order = unpack(f"<{n_triggers}I", data)
+
+        data.read(1032)  # default!
+
     return dict(
         map_id=map_id,
         difficulty_id=difficulty_id,
-        instructions=instructions
+        instructions=instructions,
+        scenario_filename=scenario_filename,
     )
 
 

tests/test_fast.py

@@ -42,6 +42,36 @@ def test_map(self):
         self.assertEqual(self.data['lobby']['seed'], -1970180596)
 
 
+class TestFastDEScenario(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        with open('tests/recs/de-50.6-scenario.aoe2record', 'rb') as handle:
+            cls.data = parse(handle)
+
+    def test_version(self):
+        self.assertEqual(self.data['version'], Version.DE)
+
+    def test_players(self):
+        players = self.data.get('players')
+        self.assertEqual(len(players), 3)
+
+
+class TestFastDEScenarioWithTriggers(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        with open('tests/recs/de-50.6-scenario-with-triggers.aoe2record', 'rb') as handle:
+            cls.data = parse(handle)
+
+    def test_version(self):
+        self.assertEqual(self.data['version'], Version.DE)
+
+    def test_players(self):
+        players = self.data.get('players')
+        self.assertEqual(len(players), 3)
+
+
 class TestFastHD(unittest.TestCase):
 
     @classmethod

tests/test_files.py

@@ -32,5 +32,10 @@ def test_files_full(self):
         parse_file_full('tests/recs/de-13.07.aoe2record')
 
     def test_files_fast(self):
+        # these files aren't supported by full header parser for now:
+        skip = {"tests/recs/de-50.6-scenario.aoe2record", "tests/recs/de-50.6-scenario-with-triggers.aoe2record"}
+
         for path in glob.glob('tests/recs/*'):
+            if path in skip:
+                continue
             parse_file_fast(path)