From 34e0de3e1f469da3817a0e124020774c653eec6a Mon Sep 17 00:00:00 2001 From: Sebastian Mohr Date: Sat, 9 Aug 2025 14:26:15 +0200 Subject: [PATCH] Added typehints and some more tests. --- beetsplug/random.py | 59 +++++++++++++++++++----------- docs/changelog.rst | 1 + test/plugins/test_random.py | 73 ++++++++++++++++++++++++++++++++++++- 3 files changed, 110 insertions(+), 23 deletions(-) diff --git a/beetsplug/random.py b/beetsplug/random.py index 330dc78e4..fcd72c83e 100644 --- a/beetsplug/random.py +++ b/beetsplug/random.py @@ -14,10 +14,14 @@ """Get a random song or album from the library.""" -import random -from itertools import groupby -from operator import attrgetter +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.ui import Subcommand, print_ @@ -69,15 +73,19 @@ class Random(BeetsPlugin): 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,28 +114,31 @@ 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 @@ -135,7 +146,11 @@ def _take_time(iter, secs, album): 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) diff --git a/docs/changelog.rst b/docs/changelog.rst index 6126f70aa..6ccdd6060 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -146,6 +146,7 @@ Other changes: - Finally removed gmusic plugin and all related code/docs as the Google Play Music service was shut down in 2020. - Updated color documentation with ``bright_*`` and ``bg_bright_*`` entries. +- Moved `beets/random.py` into `beetsplug/random.py` to cleanup core module. 2.5.1 (October 14, 2025) ------------------------ diff --git a/test/plugins/test_random.py b/test/plugins/test_random.py index cb21edf47..5b71cd126 100644 --- a/test/plugins/test_random.py +++ b/test/plugins/test_random.py @@ -20,8 +20,8 @@ from random import Random 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 @@ class RandomTest(TestHelper, unittest.TestCase): 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 == []