Added more AB controller tests and features

1 files changed, 18 insertions, 6 deletions

diff --git a/abcontroller.py b/abcontroller.py
index 8c14c0c..302bc9e 100644
--- a/abcontroller.py
+++ b/abcontroller.py
@@ -3,25 +3,34 @@ from collections import namedtuple
 _AB = namedtuple("_AB", ["a", "b"])
 
 class ABController:
-    def __init__(self):
+    def __init__(self, enabled=True):
         self.setPositionCallback = None
         self._limits = dict() # dictionary of all songs
         self._songLimits = None # list of limits for selected song
         self._currentLimits = None # a/b positions of selected limit
+        self._enabled = enabled
 
-    def setCurrentSong(self, path):
+    def _ensureSongExists(self, path):
         if path not in self._limits:
             self._limits[path] = list()
 
+    def setCurrentSong(self, path):
+        self._ensureSongExists(path)
         self._songLimits = self._limits[path]
         self._currentLimits = None
 
-    def addLimits(self, aLimit, bLimit):
-        if self._songLimits is None:
+    def addLimits(self, aLimit, bLimit, song=None):
+        if song is not None:
+            self._ensureSongExists(song)
+            songLimits = self._limits[song]
+        else:
+            songLimits = self._songLimits
+
+        if songLimits is None:
             return
 
         ab = _AB(aLimit, bLimit)
-        self._songLimits.append(ab)
+        songLimits.append(ab)
 
     def setCurrentLimits(self, index):
         if not self._songLimits:
@@ -34,5 +43,8 @@ class ABController:
         if not self._currentLimits:
             return
 
-        if position > self._currentLimits.b and self.setPositionCallback:
+        if position > self._currentLimits.b and self.setPositionCallback and self._enabled:
             self.setPositionCallback(self._currentLimits.a)
+
+    def setEnable(self, enable):
+        self._enabled = enable