Added typehints and some more tests.

Author: semohr

beetsplug/random.py

@@ -14,10 +14,14 @@
 
 """Get a random song or album from the library."""
 
+from __future__ import annotations
+
 import random
-from itertools import groupby
+from itertools import groupby, islice
 from operator import attrgetter
+from typing import Iterable, Sequence, TypeVar
 
+from beets.library import Album, Item
 from beets.plugins import BeetsPlugin
 from beets.ui import Subcommand, print_
 
@@ -69,15 +73,19 @@ def commands(self):
         return [random_cmd]
 
 
-def _length(obj, album):
+def _length(obj: Item | Album) -> float:
     """Get the duration of an item or album."""
-    if album:
+    if isinstance(obj, Album):
         return sum(i.length for i in obj.items())
     else:
         return obj.length
 
 
-def _equal_chance_permutation(objs, field="albumartist", random_gen=None):
+def _equal_chance_permutation(
+    objs: Sequence[Item | Album],
+    field: str = "albumartist",
+    random_gen: random.Random | None = None,
+) -> Iterable[Item | Album]:
     """Generate (lazily) a permutation of the objects where every group
     with equal values for `field` have an equal chance of appearing in
     any given position.
@@ -86,7 +94,7 @@ def _equal_chance_permutation(objs, field="albumartist", random_gen=None):
 
     # Group the objects by artist so we can sample from them.
     key = attrgetter(field)
-    objs.sort(key=key)
+    objs = sorted(objs, key=key)
     objs_by_artists = {}
     for artist, v in groupby(objs, key):
         objs_by_artists[artist] = list(v)
@@ -106,36 +114,43 @@ def _equal_chance_permutation(objs, field="albumartist", random_gen=None):
             del objs_by_artists[artist]
 
 
-def _take(iter, num):
+T = TypeVar("T")
+
+
+def _take(
+    iter: Iterable[T],
+    num: int,
+) -> list[T]:
     """Return a list containing the first `num` values in `iter` (or
     fewer, if the iterable ends early).
     """
-    out = []
-    for val in iter:
-        out.append(val)
-        num -= 1
-        if num <= 0:
-            break
-    return out
+    return list(islice(iter, num))
 
 
-def _take_time(iter, secs, album):
+def _take_time(
+    iter: Iterable[Item | Album],
+    secs: float,
+) -> list[Item | Album]:
     """Return a list containing the first values in `iter`, which should
     be Item or Album objects, that add up to the given amount of time in
     seconds.
     """
-    out = []
+    out: list[Item | Album] = []
     total_time = 0.0
     for obj in iter:
-        length = _length(obj, album)
+        length = _length(obj)
         if total_time + length <= secs:
             out.append(obj)
             total_time += length
     return out
 
 
 def random_objs(
-    objs, album, number=1, time=None, equal_chance=False, random_gen=None
+    objs: Sequence[Item | Album],
+    number=1,
+    time: float | None = None,
+    equal_chance: bool = False,
+    random_gen: random.Random | None = None,
 ):
     """Get a random subset of the provided `objs`.
 
@@ -152,11 +167,11 @@ def random_objs(
     if equal_chance:
         perm = _equal_chance_permutation(objs)
     else:
-        perm = objs
-        rand.shuffle(perm)  # N.B. This shuffles the original list.
+        perm = list(objs)
+        rand.shuffle(perm)
 
     # Select objects by time our count.
     if time:
-        return _take_time(perm, time * 60, album)
+        return _take_time(perm, time * 60)
     else:
         return _take(perm, number)

test/plugins/test_random.py

@@ -20,8 +20,8 @@
 
 import pytest
 
-from beets import random
 from beets.test.helper import TestHelper
+from beetsplug import random
 
 
 class RandomTest(TestHelper, unittest.TestCase):
@@ -77,3 +77,73 @@ def experiment(field, histogram=False):
         assert 0 == pytest.approx(median1, abs=1)
         assert len(self.items) // 2 == pytest.approx(median2, abs=1)
         assert stdev2 > stdev1
+
+    def test_equal_permutation_empty_input(self):
+        """Test _equal_chance_permutation with empty input."""
+        result = list(random._equal_chance_permutation([], "artist"))
+        assert result == []
+
+    def test_equal_permutation_single_item(self):
+        """Test _equal_chance_permutation with single item."""
+        result = list(random._equal_chance_permutation([self.item1], "artist"))
+        assert result == [self.item1]
+
+    def test_equal_permutation_single_artist(self):
+        """Test _equal_chance_permutation with items from one artist."""
+        items = [self.create_item(artist=self.artist1) for _ in range(5)]
+        result = list(random._equal_chance_permutation(items, "artist"))
+        assert set(result) == set(items)
+        assert len(result) == len(items)
+
+    def test_random_objs_count(self):
+        """Test random_objs with count-based selection."""
+        result = random.random_objs(
+            self.items, number=3, random_gen=self.random_gen
+        )
+        assert len(result) == 3
+        assert all(item in self.items for item in result)
+
+    def test_random_objs_time(self):
+        """Test random_objs with time-based selection."""
+        # Total length is 30 + 60 + 8*45 = 450 seconds
+        # Requesting 120 seconds should return 2-3 items
+        result = random.random_objs(
+            self.items,
+            time=2,
+            random_gen=self.random_gen,  # 2 minutes = 120 sec
+        )
+        total_time = sum(item.length for item in result)
+        assert total_time <= 120
+        # Check we got at least some items
+        assert len(result) > 0
+
+    def test_random_objs_equal_chance(self):
+        """Test random_objs with equal_chance=True."""
+
+        # With equal_chance, artist1 should appear more often in results
+        def experiment():
+            """Run the random_objs function multiple times and collect results."""
+            results = []
+            for _ in range(5000):
+                result = random.random_objs(
+                    [self.item1, self.item2],
+                    number=1,
+                    equal_chance=True,
+                    random_gen=self.random_gen,
+                )
+                results.append(result[0].artist)
+
+            return results.count(self.artist1), results.count(self.artist2)
+
+        count_artist1, count_artist2 = experiment()
+        assert abs(count_artist1 - count_artist2) < 100  # 2% deviation
+
+    def test_random_objs_empty_input(self):
+        """Test random_objs with empty input."""
+        result = random.random_objs([], number=3)
+        assert result == []
+
+    def test_random_objs_zero_number(self):
+        """Test random_objs with number=0."""
+        result = random.random_objs(self.items, number=0)
+        assert result == []