Fix unit types and add regression test.

Closes #136.

Author: GraylinKim

CHANGELOG.rst

@@ -4,10 +4,11 @@ CHANGELOG
 0.5.2 -
 --------------------
 
+* Fixed #136, unit types from tracker events are used when available.
 * Deprecated player.gateway for player.region
 * Reorganized the person/player/observer hierarchy. Top level classes are now Computer, Participant, and Observer. Participant and Computer are both children of player so any isinstance code should still work fine.
 * Player.uid now means something completely different! Use player.toon_id instead
-* Player.uid is now the user id of the player
+* Player.uid is now the user id of the player (was player.cid)
 * PersonDict can no longer be constructed from a player list and new players cannot be added by string (name). Only integer keys accepted for setting.
 * Added a sc2json script contributed by @ChrisLundquist
 * Hooked up travis-ci for continuous testing. https://travis-ci.org/GraylinKim/sc2reader

sc2reader/events/game.py

@@ -296,17 +296,27 @@ def load_context(self, replay):
 
         units = list()
         for (unit_id, unit_type, subgroup, intra_subgroup) in self.new_unit_info:
-            # Hack that defaults viking selection to fighter mode instead of assault
-            if replay.versions[1] == 2 and replay.build >= 23925 and unit_type == 71:
-                unit_type = 72
-
-            if unit_id in replay.objects:
-                unit = replay.objects[unit_id]
-                if not unit.is_type(unit_type):
-                    replay.datapack.change_type(unit, unit_type, self.frame)
+            # If we don't have access to tracker events, use selection events to create
+            # new units and track unit type changes. It won't be perfect, but it is better
+            # than nothing.
+            if not replay.tracker_events:
+                # Starting at 23925 the default viking mode is assault. Most people expect
+                # the default viking mode to be figher so fudge it a bit here.
+                if replay.versions[1] == 2 and replay.build >= 23925 and unit_type == 71:
+                    unit_type = 72
+
+                if unit_id in replay.objects:
+                    unit = replay.objects[unit_id]
+                    if not unit.is_type(unit_type):
+                        replay.datapack.change_type(unit, unit_type, self.frame)
+                else:
+                    unit = replay.datapack.create_unit(unit_id, unit_type, 0x00, self.frame)
+                    replay.objects[unit_id] = unit
+
+            # If we have tracker events, the unit must already exist and must already
+            # have the correct unit type.
             else:
-                unit = replay.datapack.create_unit(unit_id, unit_type, 0x00, self.frame)
-                replay.objects[unit_id] = unit
+                unit = replay.objects[unit_id]
 
             units.append(unit)
 

sc2reader/resources.py

@@ -229,6 +229,9 @@ def __init__(self, replay_file, filename=None, load_level=4, **options):
         self.active_units = {}
         self.game_fps = 16.0
 
+        self.tracker_events = list()
+        self.game_events = list()
+
         # Bootstrap the readers.
         self.registered_readers = defaultdict(list)
         self.register_default_readers()

test_replays/2.0.8.25604/issue136.SC2Replay

@@ -244,6 +244,14 @@ def test_cn_replays(self):
         self.assertEquals(replay.gateway, "cn")
         self.assertEquals(replay.expansion, "WoL")
 
+    def test_unit_types(self):
+        """ sc2reader#136 regression test """
+        replay = sc2reader.load_replay('test_replays/2.0.8.25604/issue136.SC2Replay')
+        hellion_times = [u.started_at for u in replay.players[0].units if u.name == 'Hellion']
+        hellbat_times = [u.started_at for u in replay.players[0].units if u.name == 'BattleHellion']
+        self.assertEquals(hellion_times, [5180, 5183])
+        self.assertEquals(hellbat_times, [6736, 6741, 7215, 7220, 12004, 12038])
+
     def test_plugins(self):
         from sc2reader.plugins.replay import APMTracker, SelectionTracker, toJSON
         factory = sc2reader.factories.SC2Factory()