Rafactored beets/random.py and moved into beetsplug/random.py

beetsplug/random.py

@@ -14,8 +14,15 @@
 
 """Get a random song or album from the library."""
 
+from __future__ import annotations
+
+import random
+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.random import random_objs
 from beets.ui import Subcommand, print_
 
 
@@ -64,3 +71,107 @@ def random_func(lib, opts, args):
 class Random(BeetsPlugin):
     def commands(self):
         return [random_cmd]
+
+
+def _length(obj: Item | Album) -> float:
+    """Get the duration of an item or album."""
+    if isinstance(obj, Album):
+        return sum(i.length for i in obj.items())
+    else:
+        return obj.length
+
+
+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.
+    """
+    rand = random_gen or random
+
+    # Group the objects by artist so we can sample from them.
+    key = attrgetter(field)
+    objs = sorted(objs, key=key)
+    objs_by_artists = {}
+    for artist, v in groupby(objs, key):
+        objs_by_artists[artist] = list(v)
+
+    # While we still have artists with music to choose from, pick one
+    # randomly and pick a track from that artist.
+    while objs_by_artists:
+        # Choose an artist and an object for that artist, removing
+        # this choice from the pool.
+        artist = rand.choice(list(objs_by_artists.keys()))
+        objs_from_artist = objs_by_artists[artist]
+        i = rand.randint(0, len(objs_from_artist) - 1)
+        yield objs_from_artist.pop(i)
+
+        # Remove the artist if we've used up all of its objects.
+        if not objs_from_artist:
+            del objs_by_artists[artist]
+
+
+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).
+    """
+    return list(islice(iter, num))
+
+
+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: list[Item | Album] = []
+    total_time = 0.0
+    for obj in iter:
+        length = _length(obj)
+        if total_time + length <= secs:
+            out.append(obj)
+            total_time += length
+    return out
+
+
+def random_objs(
+    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`.
+
+    If `number` is provided, produce that many matches. Otherwise, if
+    `time` is provided, instead select a list whose total time is close
+    to that number of minutes. If `equal_chance` is true, give each
+    artist an equal chance of being included so that artists with more
+    songs are not represented disproportionately.
+    """
+    rand = random_gen or random
+
+    # Permute the objects either in a straightforward way or an
+    # artist-balanced way.
+    if equal_chance:
+        perm = _equal_chance_permutation(objs)
+    else:
+        perm = list(objs)
+        rand.shuffle(perm)
+
+    # Select objects by time our count.
+    if time:
+        return _take_time(perm, time * 60)
+    else:
+        return _take(perm, number)

docs/changelog.rst

@@ -107,6 +107,7 @@ Other changes:
 * Refactored library.py file by splitting it into multiple modules within the
   beets/library directory.
 * Added a test to check that all plugins can be imported without errors.
+* Moved `beets/random.py` into `beetsplug/random.py` to cleanup core module.
 
 2.3.1 (May 14, 2025)
 --------------------

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,74 @@ 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 ratio
+            return results.count(self.artist1), results.count(self.artist2)
+
+        count_artist1, count_artist2 = experiment()
+        assert 1 - count_artist1 / count_artist2 < 0.1  # 10% 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 == []