Properly handle hallucinated units. Fixes #153.

Author: GraylinKim

sc2reader/data/__init__.py

@@ -37,7 +37,7 @@ class Unit(object):
     Represents an in-game unit.
     """
 
-    def __init__(self, unit_id, flags):
+    def __init__(self, unit_id):
         #: A reference to the player that currently owns this unit. Only available for 2.0.8+ replays.
         self.owner = None
 
@@ -66,8 +66,6 @@ def __init__(self, unit_id, flags):
         #: behind the fog of war is targetted.
         self.id = unit_id
 
-        self.flags = flags
-
         #: A reference to the unit type this unit is current in.
         #: e.g. SeigeTank is a different type than SeigeTankSeiged
         self._type_class = None
@@ -76,8 +74,14 @@ def __init__(self, unit_id, flags):
         #: in order by frame the type was acquired.
         self.type_history = OrderedDict()
 
-        #: Is this unit type a hallucinated one? Unsure of this flag..
-        self.hallucinated = (flags & 2 == 2)
+        #: Is this unit type a hallucinated one?
+        self.hallucinated = False
+
+        self.flags = 0
+
+    def apply_flags(self, flags):
+        self.flags = flags
+        self.hallucinated = flags & 2 == 2
 
     def set_type(self, unit_type, frame):
         self._type_class = unit_type
@@ -231,14 +235,14 @@ def __init__(self, build_id):
         #: A dictionary mapping integer ids to available abilities.
         self.abilities = dict()
 
-    def create_unit(self, unit_id, unit_type, unit_flags, frame):
+    def create_unit(self, unit_id, unit_type, frame):
         """
         :param unit_id: The unique id of this unit.
         :param unit_type: The unit type to assign to the new unit
 
         Creates a new unit and assigns it to the specified type.
         """
-        unit = Unit(unit_id, unit_flags)
+        unit = Unit(unit_id)
         self.change_type(unit, unit_type, frame)
         return unit
 

sc2reader/engine/plugins/context.py

@@ -55,7 +55,7 @@ def handleTargetAbilityEvent(self, event, replay):
             # Often when the target_unit_id is not in replay.objects it is 0 because it
             # is a target building/destructable hidden by fog of war. Perhaps we can match
             # it through the fog using location?
-            unit = replay.datapack.create_unit(event.target_unit_id, event.target_unit_type, 0x00, event.frame)
+            unit = replay.datapack.create_unit(event.target_unit_id, event.target_unit_type, event.frame)
             event.target = unit
             replay.objects[event.target_unit_id] = unit
 
@@ -64,7 +64,8 @@ def handleSelectionEvent(self, event, replay):
             return
 
         units = list()
-        for (unit_id, unit_type, subgroup, intra_subgroup) in event.new_unit_info:
+        # TODO: Blizzard calls these subgroup flags but that doesn't make sense right now
+        for (unit_id, unit_type, subgroup_flags, intra_subgroup_flags) in event.new_unit_info:
             # 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.
@@ -79,14 +80,17 @@ def handleSelectionEvent(self, event, replay):
                     if not unit.is_type(unit_type):
                         replay.datapack.change_type(unit, unit_type, event.frame)
                 else:
-                    unit = replay.datapack.create_unit(unit_id, unit_type, 0x00, event.frame)
+                    unit = replay.datapack.create_unit(unit_id, unit_type, event.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.objects[unit_id]
 
+            # Selection events hold flags on units (like hallucination)
+            unit.apply_flags(intra_subgroup_flags)
+
             units.append(unit)
 
         event.new_units = event.objects = units
@@ -110,7 +114,7 @@ def handleUnitBornEvent(self, event, replay):
             event.unit = replay.objects[event.unit_id]
         else:
             # TODO: How to tell if something is hallucination?
-            event.unit = replay.datapack.create_unit(event.unit_id, event.unit_type_name, 0, event.frame)
+            event.unit = replay.datapack.create_unit(event.unit_id, event.unit_type_name, event.frame)
             replay.objects[event.unit_id] = event.unit
 
         replay.active_units[event.unit_id_index] = event.unit
@@ -189,7 +193,7 @@ def handleUnitInitEvent(self, event, replay):
             event.unit = replay.objects[event.unit_id]
         else:
             # TODO: How to tell if something is hallucination?
-            event.unit = replay.datapack.create_unit(event.unit_id, event.unit_type_name, 0, event.frame)
+            event.unit = replay.datapack.create_unit(event.unit_id, event.unit_type_name, event.frame)
             replay.objects[event.unit_id] = event.unit
 
         replay.active_units[event.unit_id_index] = event.unit