trackdb: Give state functions to the TrackAllocator

So that we don't try to pickle out things that can't be pickled

Signed-off-by: Anna Schumaker <Anna@NoWheyCreamery.com>

trackdb/allocator.py

@@ -23,6 +23,17 @@ class TrackAllocator:
         with self.lock:
             return self.tracks.get(id, None)
 
+    def __getstate__(self):
+        with self.lock:
+            return { "tracks" : self.tracks,
+                     "nextid" : self.nextid }
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self.lock = threading.Lock()
+        self.Added = publisher.Publisher()
+        self.Removed = publisher.Publisher()
+
     def allocate(self, lib, filepath):
         try:
             return self.__alloc_track__(lib, filepath)

trackdb/test_allocator.py

@@ -58,3 +58,20 @@ class TestTrackAllocator(unittest.TestCase):
         self.assertEqual(alloc.tracks, { })
         self.assertEqual(alloc.Added.subscribers, set())
         self.assertEqual(alloc.Removed.subscribers, set())
+
+    def test_allocator_state(self):
+        lib = library.LibraryPath(test_tracks)
+        alloc = allocator.TrackAllocator()
+        track = alloc.allocate(lib, test_tracks / "01 - Test Track.ogg")
+
+        state = alloc.__getstate__()
+        self.assertEqual(state, { "tracks" : { 0 : track }, 
+                                  "nextid" : 1 })
+
+        alloc.__dict__.clear()
+        alloc.__setstate__(state)
+        self.assertEqual(alloc.tracks, { 0 : track })
+        self.assertEqual(alloc.nextid, 1)
+        self.assertIsInstance(alloc.lock, type(threading.Lock()))
+        self.assertIsInstance(alloc.Added, publisher.Publisher)
+        self.assertIsInstance(alloc.Removed, publisher.Publisher)