Improved the tracker events processing.

Tracker Events now provide:

	unit.owner
	player.units
	unit.started_at
	unit.finished_at
	unit.died_at

Author: GraylinKim

sc2reader/data/__init__.py

@@ -26,6 +26,18 @@
 class Unit(object):
 
     def __init__(self, unit_id, flags):
+        #: A reference to the player that owns this unit
+        self.owner = None
+
+        #: The frame the unit was started at
+        self.started_at = None
+
+        #: The frame the unit was finished at
+        self.finished_at = None
+
+        #: The frame the unit died at
+        self.died_at = None
+
         self.id = unit_id
         self.flags = flags
         self._type_class = None

sc2reader/events/base.py

@@ -19,6 +19,7 @@ def load_context(self, replay):
         if replay.versions[1]==1 or (replay.versions[1]==2 and replay.build < 24247):
             if self.pid <= len(replay.people):
                 self.player = replay.person[self.pid]
+                self.player.events.append(self)
             elif self.pid != 16:
                 self.logger.error("Bad pid ({0}) for event {1} at {2}.".format(self.pid, self.__class__, Length(seconds=self.second)))
             else:
@@ -27,6 +28,7 @@ def load_context(self, replay):
         else:
             if self.pid < len(replay.clients):
                 self.player = replay.client[self.pid]
+                self.player.events.append(self)
             elif self.pid != 16:
                 self.logger.error("Bad pid ({0}) for event {1} at {2}.".format(self.pid, self.__class__, Length(seconds=self.second)))
             else:

sc2reader/events/tracker.py

@@ -1,18 +1,20 @@
 # -*- coding: utf-8 -*-
 import json
+from sc2reader.utils import Length
 
 class TrackerEvent(object):
     def __init__(self, frames):
         #: The frame of the game this event was applied
         self.frame = frames
 
-    def __str__(self):
-        """ Dumps the event data to a json string. """
-        return json.dumps(self.__dict__)
-
     def load_context(self, replay):
         pass
 
+    def _str_prefix(self):
+        return "{0}\t ".format(Length(seconds=int(self.frame/16)))
+
+    def __str__(self):
+        return self._str_prefix() + self.name
 
 class PlayerStatsEvent(TrackerEvent):
     name = 'PlayerStatsEvent'
@@ -131,6 +133,8 @@ def __init__(self, frames, data):
     def load_context(self, replay):
         self.player = replay.player[self.pid]
 
+    def __str__(self):
+        return self._str_prefix()+"{0: >15} - Stats Update".format(self.player)
 
 class UnitBornEvent(TrackerEvent):
     name = 'UnitBornEvent'
@@ -195,9 +199,15 @@ def load_context(self, replay):
 
         replay.active_units[self.unit_id_index] = self.unit
         self.unit.location = self.location
-        self.unit.birth = self.frame
-        self.unit.owner = self.unit_upkeeper
+        self.unit.started_at = self.frame
+        self.unit.finished_at = self.frame
+
+        if self.unit_upkeeper:
+            self.unit.owner = self.unit_upkeeper
+            self.unit.owner.units.append(self.unit)
 
+    def __str__(self):
+        return self._str_prefix()+"{0: >15} - Unit born {1}".format(self.unit_upkeeper,self.unit)
 
 class UnitDiedEvent(TrackerEvent):
     name = 'UnitDiedEvent'
@@ -235,7 +245,7 @@ def __init__(self, frames, data):
     def load_context(self, replay):
         if self.unit_id in replay.objects:
             self.unit = replay.objects[self.unit_id]
-            self.unit.death = self.frame
+            self.unit.died_at = self.frame
             self.unit.location = self.location
             if self.unit_id_index in replay.active_units:
                 del replay.active_units[self.unit_id_index]
@@ -249,6 +259,9 @@ def load_context(self, replay):
         elif self.killer_pid:
             pass#print "Unknown killer pid", self.killer_pid
 
+    def __str__(self):
+        return self._str_prefix()+"{0: >15} - Unit died {1}.".format(self.unit.owner, self.unit)
+
 
 class UnitOwnerChangeEvent(TrackerEvent):
     name = 'UnitOwnerChangeEvent'
@@ -293,10 +306,19 @@ def load_context(self, replay):
 
         if self.unit_id in replay.objects:
             self.unit = replay.objects[self.unit_id]
-            self.unit.owner = self.unit_upkeeper
         else:
             print "Unit owner changed before it was born!"
 
+        if self.unit_upkeeper:
+            if unit.owner:
+                print "stduff"
+                unit.owner.units.remove(unit)
+            unit.owner = self.unit_upkeeper
+            self.unit_upkeeper.units.append(unit)
+
+    def __str__(self):
+        return self._str_prefix()+"{0: >15} took {1}".format(self.unit_upkeeper, self.unit)
+
 
 class UnitTypeChangeEvent(TrackerEvent):
     name = 'UnitTypeChangeEvent'
@@ -326,6 +348,9 @@ def load_context(self, replay):
         else:
             print "Unit type changed before it was born!"
 
+    def __str__(self):
+        return self._str_prefix()+"{0: >15} - Unit {0} type changed to {1}".format(self.unit.owner, self.unit, self.unit_type_name)
+
 
 class UpgradeCompleteEvent(TrackerEvent):
     name = 'UpgradeCompleteEvent'
@@ -354,6 +379,8 @@ def load_context(self, replay):
         # TODO: We don't have upgrade -> ability maps
         # TODO: we can probably do the same thing we did for units
 
+    def __str__(self):
+        return self._str_prefix()+"{0: >15} - {1}upgrade completed".format(self.player, self.upgrade_type_name)
 
 class UnitInitEvent(TrackerEvent):
     name = 'UnitInitEvent'
@@ -419,9 +446,16 @@ def load_context(self, replay):
             replay.objects[self.unit_id] = self.unit
 
         replay.active_units[self.unit_id_index] = self.unit
-        self.unit.owner = self.unit_upkeeper
         self.unit.location = self.location
-        self.unit.birth = self.frame
+        self.unit.started_at = self.frame
+        # self.unit.finished_at = self.frame
+
+        if self.unit_upkeeper:
+            self.unit.owner = self.unit_upkeeper
+            self.unit.owner.units.append(self.unit)
+
+    def __str__(self):
+        return self._str_prefix()+"{0: >15} - Unit inititated {1}".format(self.unit_upkeeper, self.unit)
 
 
 class UnitDoneEvent(TrackerEvent):
@@ -445,9 +479,12 @@ def __init__(self, frames, data):
     def load_context(self, replay):
         if self.unit_id in replay.objects:
             self.unit = replay.objects[self.unit_id]
+            self.unit.finished_at = self.frame
         else:
             print "Unit done before it was started!"
 
+    def __str__(self):
+        return self._str_prefix()+"{0: >15} - Unit {1} done".format(self.unit.owner, self.unit)
 
 class UnitPositionsEvent(TrackerEvent):
     name = 'UnitPositionsEvent'
@@ -482,3 +519,6 @@ def load_context(self, replay):
                 self.units.append(unit)
             else:
                 print "Unit moved that doesn't exist!"
+
+    def __str__(self):
+        return self._str_prefix()+"Unit positions update"

sc2reader/objects.py

@@ -201,6 +201,10 @@ def __init__(self, pid, name):
         super(Player,self).__init__(pid, name)
         self.is_observer = False
 
+        #: A list of references to the units the player had this game
+        self.units = list()
+
+
     @property
     def url(self):
         """The player's battle.net profile url"""

sc2reader/plugins/replay.py

@@ -115,7 +115,6 @@ def SelectionTracker(replay):
         player_selections = GameState(PlayerSelection())
         for event in person.events:
             error = False
-
             if event.name == 'SelectionEvent':
                 selections = player_selections[event.frame]
                 control_group = selections[event.control_group].copy()

sc2reader/resources.py

@@ -281,6 +281,7 @@ def __init__(self, replay_file, filename=None, load_level=4, **options):
         for event in self.events:
             event.load_context(self)
 
+
     def load_details(self):
         if 'replay.attributes.events' in self.raw_data:
             # Organize the attribute data to be useful
@@ -475,21 +476,17 @@ def load_messages(self):
     def load_events(self):
         # Copy the events over
         # TODO: the events need to be fixed both on the reader and processor side
-        if 'replay.game.events' in self.raw_data:
-            self.game_events = self.raw_data['replay.game.events']
-            self.events = sorted(self.game_events+self.events, key=lambda e: e.frame)
+        if 'replay.game.events' not in self.raw_data:
+            return
+
+        self.game_events = self.raw_data['replay.game.events']
+        self.events = sorted(self.game_events+self.events, key=lambda e: e.frame)
 
         # hideous hack for HotS 2.0.0.23925, see https://github.com/GraylinKim/sc2reader/issues/87
         if self.events and self.events[-1].frame > self.frames:
             self.frames = self.events[-1].frame
             self.length = utils.Length(seconds=int(self.frames/self.game_fps))
 
-        self.camera_events = list()
-        self.selection_events = list()
-        self.ability_events = list()
-        for event in self.events:
-            if event.pid != 16 and hasattr(event,'player'):
-                event.player.events.append(event)
 
     def load_tracker_events(self):
         if 'replay.tracker.events' not in self.raw_data: