Improve SelectionTracker processing speed.

Just a combination of improvements using stdlibrary tools.

Author: GraylinKim

sc2reader/plugins/utils.py

@@ -3,9 +3,11 @@
 import json
 from functools import wraps
 from datetime import datetime
-from sc2reader import log_utils
+from sc2reader.log_utils import loggable
 
 from functools import wraps
+from bisect import bisect_left
+
 def plugin(func):
     @wraps(func)
     def wrapper(**options):
@@ -26,28 +28,45 @@ def default(self, obj):
 
 class GameState(dict):
     def __init__(self, initial_state):
+        self._frames = list()
+        self._frameset = set()
         self[0] = initial_state
 
     def __getitem__(self, frame):
         if frame in self:
             return super(GameState, self).__getitem__(frame)
 
-        # If this particular key doesn't exist, create it
-        prev_frame = sorted(filter(lambda key: key < frame, self.keys()))[-1]
+        # Get the previous frame from our sorted frame list
+        # bisect_left returns the left most key where an item is
+        # less than or equal to the value in that key. If it is
+        # less than we need to subtract 1
+        key = bisect_left(self._frames, frame)
+        if key == len(self._frames) or self._frames[key] > frame:
+            prev_frame = self._frames[key - 1]
+        else:
+            prev_frame = self._frames[key]
+
+        # Copy the previous state and use it as our basis here
         state = self[prev_frame].copy()
         self[frame] = state
         return state
 
+    def __setitem__(self, frame, value):
+        if frame not in self._frameset:
+            self._frames.insert(bisect_left(self._frames, frame), frame)
+            self._frameset.add(frame)
+
+        super(GameState, self).__setitem__(frame, value)
+
 
+@loggable
 class UnitSelection(object):
-    def __init__(self, *new_objects):
-        self.logger = log_utils.get_logger(UnitSelection)
-        self.objects = list()
-        self.select(new_objects)
+    def __init__(self, objects=None):
+        self.objects = objects or list()
 
     def select(self, new_objects):
-        object_set = set(self.objects+list(new_objects))
-        self.objects = sorted(object_set, key=lambda obj: obj.id)
+        new_set = set(self.objects+list(new_objects))
+        self.objects = sorted(new_set,key=lambda obj: obj.id)
 
     def deselect(self, mode, data):
         if mode == 0x01:
@@ -71,7 +90,7 @@ def __str__(self):
         return ', '.join(str(obj) for obj in self.objects)
 
     def copy(self):
-        return UnitSelection(*self.objects)
+        return UnitSelection(self.objects[:])
 
 
 class PlayerSelection(dict):
@@ -82,5 +101,5 @@ def __init__(self):
     def copy(self):
         new = PlayerSelection()
         for bank, selection in self.iteritems():
-            new[bank] = UnitSelection(*selection.objects)
+            new[bank] = UnitSelection(selection.objects[:])
         return new