Merge pull request #46 from r-Norge/feature/queue_type_hints

Feature/queue type hints

.github/workflows/build.yml

@@ -24,6 +24,12 @@ jobs:
           pip install isort
           isort --check --sp setup.cfg ${CheckFiles}
 
+      - name: Run tests
+        run: |
+          pip install pytest
+          pip install -r requirements.txt
+          pytest
+
   build-and-push:
     needs: flake8
     runs-on: ubuntu-latest

musicbot/utils/mixplayer/mixqueue.py

@@ -1,13 +1,14 @@
-# Discord Packages
-from lavalink import AudioTrack
-
 from collections import OrderedDict, deque
 from itertools import chain, cycle, islice
 from random import shuffle
-from typing import List, Union
+from typing import Generic, Iterable, Iterator, List, Tuple, TypeVar, Union
+
+# Would like to ensure the T has a "requester" attribute, but don't know if that is possible
+T = TypeVar('T')
+QueueType = List[T]
 
 
-def roundrobin(*iterables):
+def roundrobin(*iterables: Iterable[T]) -> Iterator[T]:
     """roundrobin('ABC', 'D', 'EF') --> A D E B F C"""
     # Recipe credited to George Sakkis
     num_active = len(iterables)
@@ -21,61 +22,62 @@ def roundrobin(*iterables):
             nexts = cycle(islice(nexts, num_active))
 
 
-class MixQueue:
+class MixQueue(Generic[T]):
     def __init__(self):
-        self.queues: OrderedDict[int, List[AudioTrack]] = OrderedDict()
-        self.priority_queue = []
+        self.queues: OrderedDict[int, QueueType] = OrderedDict()
+        self.priority_queue: QueueType = []
         self._history = deque(maxlen=11)  # 10 + current
         self.looping = False
         self.loop_offset = 0
 
-    def __str__(self):
+    def __str__(self) -> str:
         tmp = ''
         for requester in self.queues.keys():
             user_queue = list(map(str, self.queues[requester]))
             tmp += f'{requester}: {user_queue}\n'
         return tmp
 
-    def __bool__(self):
+    def __bool__(self) -> bool:
         return not self.empty
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[T]:
         global_queue = roundrobin(*[x for x in self.queues.values()])
         priority_queue = self.priority_queue
         out = chain(priority_queue, global_queue)
         return out
 
-    def __len__(self):
+    def __len__(self) -> int:
         length = 0
         for _ in self:
             length += 1
         return length
 
-    def get_queue(self):
+    def get_queue(self) -> QueueType:
         if self.looping:
             offset = self.loop_offset
             return list(self)[offset:] + list(self)[:offset]
         else:
             return list(self)
 
-    def clear(self):
+    def clear(self) -> None:
         self.queues = OrderedDict()
         self.priority_queue = []
         self.looping = False
         self.loop_offset = 0
 
     # if dual is true also returns global positions of tracks
-    def get_user_queue(self, requester: int, dual: bool = False):
-        queue = self.queues.get(requester, [])
-        if dual and queue:
-            pos = [self._loc_to_glob(requester, i) for i in range(len(queue))]
-            if self.looping:
-                pos = [(p + self.loop_offset) % len(self) for p in pos]
-            combined = zip(queue, pos)
-            return list(combined)
-        return queue
+    def get_user_queue(self, requester: int) -> QueueType:
+        return self.queues.get(requester, [])
 
-    def pop_first(self):
+    def get_user_queue_with_index(self, requester: int) -> List[Tuple[T, int]]:
+        queue = self.queues.get(requester, [])
+        pos = [self._loc_to_glob(requester, i) for i in range(len(queue))]
+        if self.looping:
+            pos = [(p + self.loop_offset) % len(self) for p in pos]
+        combined = zip(queue, pos)
+        return list(combined)
+
+    def pop_first(self) -> Union[T, None]:
         if self.priority_queue:
             next_track = self.priority_queue.pop(0)
             self._history.append(next_track)
@@ -96,15 +98,16 @@ class MixQueue:
                     print(len(self))
                     self.loop_offset = min(self.loop_offset, len(self)-1)
             else:
-                next_track = self.queues[self.first_queue].pop(0)
-                self._shuffle()
-                self._clear_empty()
-                self._history.append(next_track)
-                return next_track
+                if first_queue := self.first_queue:
+                    next_track = self.queues[first_queue].pop(0)
+                    self._shift_queues()
+                    self._clear_empty()
+                    self._history.append(next_track)
+                    return next_track
         except KeyError:
             pass
 
-    def add_track(self, requester: int, track: AudioTrack, pos: int = None):
+    def add_track(self, requester: int, track: T, pos: Union[int, None] = None) -> Tuple[T, int, int]:
         user_queue = self.queues.get(requester)
         if user_queue is None:
             self.queues[requester] = [track]
@@ -119,20 +122,20 @@ class MixQueue:
         # Return info about track position
         return track, self._loc_to_glob(requester, localpos), localpos
 
-    def add_next_track(self, track: AudioTrack):
+    def add_next_track(self, track: T) -> None:
         self.priority_queue.append(track)
 
-    def remove_user_queue(self, requester: int):
+    def remove_user_queue(self, requester: int) -> None:
         user_queue = self.queues.get(requester, [])
         if user_queue:
             if self.looping:
                 # Should work, kinda hacky.
-                for pos, track in reversed(list(enumerate(user_queue))):
+                for pos, _ in reversed(list(enumerate(user_queue))):
                     self.remove_user_track(requester, pos)
             else:
                 self.queues.pop(requester)
 
-    def remove_user_track(self, requester: int, pos: int) -> Union[None, AudioTrack]:
+    def remove_user_track(self, requester: int, pos: int) -> Union[None, T]:
         user_queue = self.queues.get(requester)
         if user_queue is not None:
             if pos < len(user_queue):
@@ -143,22 +146,22 @@ class MixQueue:
                 self._clear_empty()
                 return track
 
-    def remove_global_track(self, pos: int):
+    def remove_global_track(self, pos: int) -> Union[T, None]:
         # Get the actual index of the song
-        pos = list(self).index(self.get_queue()[pos])
+        index = list(self).index(self.get_queue()[pos])
 
-        if pos <= self.loop_offset:
+        if index <= self.loop_offset:
             self.loop_offset -= 1
 
-        q, pos = self._glob_to_loc(pos)
-        if q is None or pos is None:
+        q, index = self._glob_to_loc(index)
+        if q is None or index is None:
             return
-        queue = self.queues.get(q)
-        track = queue.pop(pos)
-        self._clear_empty()
-        return track
+        if queue := self.queues.get(q):
+            track = queue.pop(index)
+            self._clear_empty()
+            return track
 
-    def move_user_track(self, requester: int, initial: int, final: int):
+    def move_user_track(self, requester: int, initial: int, final: int) -> None:
         queue = self.queues.get(requester, [])
         if queue:
             try:
@@ -168,12 +171,12 @@ class MixQueue:
             except IndexError:
                 pass
 
-    def shuffle_user_queue(self, requester: int):
+    def shuffle_user_queue(self, requester: int) -> None:
         queue = self.queues.get(requester, [])
         if queue:
             shuffle(queue)
 
-    def enable_looping(self, looping):
+    def enable_looping(self, looping: bool) -> None:
         if (not self.looping) and looping:  # Not enabled and turning on
             self.looping = looping
             self.loop_offset = 0
@@ -187,16 +190,17 @@ class MixQueue:
             self.loop_offset = 0
 
     # Switches the order of user queues
-    def _shuffle(self):
-        self.queues.move_to_end(self.first_queue)
+    def _shift_queues(self) -> None:
+        if idx := self.first_queue:
+            self.queues.move_to_end(idx)
 
     # Removes empty user queues
-    def _clear_empty(self):
+    def _clear_empty(self) -> None:
         to_remove = [q for q in reversed(self.queues) if not self.queues[q]]
         for i in to_remove:
             self.queues.pop(i)
 
-    def _loc_to_glob(self, requester, pos):
+    def _loc_to_glob(self, requester, pos) -> int:
         globpos = len(self.priority_queue)
         passed = False
         for key, queue in self.queues.items():
@@ -210,7 +214,7 @@ class MixQueue:
                     globpos += 1
         return globpos
 
-    def _glob_to_loc(self, pos: int):
+    def _glob_to_loc(self, pos: int) -> Tuple[Union[T, None], Union[int, None]]:
         try:
             song = next(islice(self, pos, pos + 1))
         except (ValueError, StopIteration):
@@ -224,18 +228,19 @@ class MixQueue:
         for i, s in enumerate(user_queue):
             if s == song:
                 return song.requester, i
+        return None, None
 
     @property
-    def first_queue(self):
+    def first_queue(self) -> Union[int, None]:
         try:
             return list(self.queues.keys())[0]
         except IndexError:
             pass
 
     @property
-    def empty(self):
+    def empty(self) -> bool:
         return len(self) == 0
 
     @property
-    def history(self):
+    def history(self) -> QueueType:
         return list(reversed(self._history))

musicbot/utils/mixplayer/player.py

@@ -65,7 +65,9 @@ class MixPlayer(DefaultPlayer):
         self.queue.shuffle_user_queue(requester)
 
     def user_queue(self, user: int, dual: bool = False):
-        return self.queue.get_user_queue(user, dual)
+        if dual:
+            return self.queue.get_user_queue_with_index(user)
+        return self.queue.get_user_queue(user)
 
     def global_queue(self):
         return self.queue.get_queue()

musicbot/utils/mixplayer/test_queue.py

@@ -0,0 +1,123 @@
+from .mixqueue import MixQueue
+
+
+class TrackMock:
+    def __init__(self, requester, title):
+        self.requester = requester
+        self.title = str(requester)+title
+
+    def __eq__(self, other):
+        return self.title == other.title
+
+    def __str__(self):
+        return f"({self.requester}: {self.title})"
+
+    def __repr__(self):
+        return str(self)
+
+
+class TestMixQueue():
+    def list_to_requests(self, list_of_str):
+        tracks = []
+        for item in list_of_str:
+            requester = int(item[0])
+            title = str(item[1:])
+            tracks.append(TrackMock(requester, title))
+        return tracks
+
+    def setup_baisc_queue(self):
+        queue = MixQueue()
+        queue.add_track(2, TrackMock(2, "a"))
+        queue.add_track(2, TrackMock(2, "b"))
+
+        queue.add_track(1, TrackMock(1, "a"))
+        queue.add_track(1, TrackMock(1, "b"))
+        queue.add_track(1, TrackMock(1, "c"))
+        queue.add_track(1, TrackMock(1, "d"))
+
+        queue.add_track(3, TrackMock(3, "a"))
+        queue.add_track(3, TrackMock(3, "b"))
+        return queue
+
+    def test_queue_mixing_double(self):
+        queue = MixQueue()
+        requester1 = 1
+        requester2 = 2
+        queue.add_track(requester1, TrackMock(1, "a"))
+        queue.add_track(requester1, TrackMock(1, "b"))
+        queue.add_track(requester1, TrackMock(1, "c"))
+        queue.add_track(requester1, TrackMock(1, "d"))
+        queue.add_track(requester1, TrackMock(1, "e"))
+
+        queue.add_track(requester2, TrackMock(2, "a"))
+        queue.add_track(requester2, TrackMock(2, "b"))
+        queue.add_track(requester2, TrackMock(2, "c"))
+
+        assert list(queue) == self.list_to_requests(["1a", "2a", "1b", "2b", "1c", "2c", "1d", "1e"])
+
+    def test_queue_mixing_triple(self):
+        queue = self.setup_baisc_queue()
+        assert list(queue) == self.list_to_requests(["2a", "1a", "3a", "2b", "1b", "3b", "1c", "1d"])
+
+    def test_priority_queue(self):
+        queue = self.setup_baisc_queue()
+
+        queue.add_next_track("priority_item1")
+        queue.add_next_track("priority_item2")
+
+        assert queue.pop_first() == "priority_item1"
+        assert queue.pop_first() == "priority_item2"
+        assert queue.pop_first() == TrackMock(2, "a")
+        assert queue.pop_first() == TrackMock(1, "a")
+
+    def test_user_queue(self):
+        queue = self.setup_baisc_queue()
+        assert queue.get_user_queue(1) == [TrackMock(1, "a"), TrackMock(1, "b"), TrackMock(1, "c"), TrackMock(1, "d")]
+        assert queue.get_user_queue(2) == [TrackMock(2, "a"), TrackMock(2, "b")]
+        assert queue.get_user_queue(3) == [TrackMock(3, "a"), TrackMock(3, "b")]
+
+    def test_user_queue_with_index(self):
+        queue = self.setup_baisc_queue()
+        assert queue.get_user_queue_with_index(1) == [(TrackMock(1, "a"), 1), (TrackMock(1, "b"), 4),
+                                                      (TrackMock(1, "c"), 6), (TrackMock(1, "d"), 7)]
+        assert queue.get_user_queue_with_index(2) == [(TrackMock(2, "a"), 0), (TrackMock(2, "b"), 3)]
+        assert queue.get_user_queue_with_index(3) == [(TrackMock(3, "a"), 2), (TrackMock(3, "b"), 5)]
+
+    def test_looping_behavior(self):
+        queue = self.setup_baisc_queue()
+
+        # Before looping
+        assert queue.get_user_queue_with_index(1) == [(TrackMock(1, "a"), 1), (TrackMock(1, "b"), 4),
+                                                      (TrackMock(1, "c"), 6), (TrackMock(1, "d"), 7)]
+        assert queue.get_user_queue_with_index(2) == [(TrackMock(2, "a"), 0), (TrackMock(2, "b"), 3)]
+        assert queue.get_user_queue_with_index(3) == [(TrackMock(3, "a"), 2), (TrackMock(3, "b"), 5)]
+        assert list(queue) == self.list_to_requests(["2a", "1a", "3a", "2b", "1b", "3b", "1c", "1d"])
+
+        queue.enable_looping(True)
+        assert queue.pop_first() == TrackMock(2, "a")
+        assert list(queue) == self.list_to_requests(["2a", "1a", "3a", "2b", "1b", "3b", "1c", "1d"])
+        assert queue.get_user_queue_with_index(1) == [(TrackMock(1, "a"), 2), (TrackMock(1, "b"), 5),
+                                                      (TrackMock(1, "c"), 7), (TrackMock(1, "d"), 0)]
+        assert queue.get_user_queue_with_index(2) == [(TrackMock(2, "a"), 1), (TrackMock(2, "b"), 4)]
+        assert queue.get_user_queue_with_index(3) == [(TrackMock(3, "a"), 3), (TrackMock(3, "b"), 6)]
+
+        popped_queue_before_loop_point = self.list_to_requests(["1a", "3a", "2b", "1b", "3b", "1c", "1d"])
+        for item in popped_queue_before_loop_point:
+            assert item == queue.pop_first()
+
+        popped_queue_after_loop_point = self.list_to_requests(["2a", "1a", "3a", "2b", "1b", "3b", "1c", "1d"])
+        for item in popped_queue_after_loop_point:
+            assert item == queue.pop_first()
+
+        popped_queue_after_loop_point = self.list_to_requests(["2a", "1a", "3a", "2b"])
+        for item in popped_queue_after_loop_point:
+            assert item == queue.pop_first()
+
+        queue.enable_looping(False)
+        # When the queue is disabled we take the current requester order of
+        # the looped queue and regenerate the mixed queue from that
+        popped_queue_after_loop_disabled = self.list_to_requests(["1b", "3b", "2a", "1c", "3a", "2b", "1d", "1a"])
+        for item in popped_queue_after_loop_disabled:
+            assert item == queue.pop_first()
+
+        assert list(queue) == []