Overall refactor of random plugin. Added length property to albums.

Author: semohr

beets/library/models.py

@@ -616,6 +616,11 @@ def try_sync(self, write, move, inherit=True):
         for item in self.items():
             item.try_sync(write, move)
 
+    @property
+    def length(self) -> float:
+        """Return the total length of all items in this album in seconds."""
+        return sum(item.length for item in self.items())
+
 
 class Item(LibModel):
     """Represent a song or track."""

beetsplug/random.py

@@ -1,42 +1,34 @@
-# This file is part of beets.
-# Copyright 2016, Philippe Mongeau.
-#
-# Permission is hereby granted, free of charge, to any person obtaining
-# a copy of this software and associated documentation files (the
-# "Software"), to deal in the Software without restriction, including
-# without limitation the rights to use, copy, modify, merge, publish,
-# distribute, sublicense, and/or sell copies of the Software, and to
-# permit persons to whom the Software is furnished to do so, subject to
-# the following conditions:
-#
-# The above copyright notice and this permission notice shall be
-# included in all copies or substantial portions of the Software.
-
 """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 typing import TYPE_CHECKING, Any, Iterable, Sequence, TypeVar
 
-from beets.library import Album, Item
 from beets.plugins import BeetsPlugin
 from beets.ui import Subcommand, print_
 
+if TYPE_CHECKING:
+    import optparse
+
+    from beets.library import Album, Item, Library
+
+    T = TypeVar("T", bound=Item | Album)
+
 
-def random_func(lib, opts, args):
+def random_func(lib: Library, opts: optparse.Values, args: list[str]):
     """Select some random items or albums and print the results."""
     # Fetch all the objects matching the query into a list.
-    if opts.album:
-        objs = list(lib.albums(args))
-    else:
-        objs = list(lib.items(args))
+    objs = lib.albums(args) if opts.album else lib.items(args)
 
     # Print a random subset.
     objs = random_objs(
-        objs, opts.album, opts.number, opts.time, opts.equal_chance
+        objs=list(objs),
+        number=opts.number,
+        time_minutes=opts.time,
+        equal_chance=opts.equal_chance,
     )
     for obj in objs:
         print_(format(obj))
@@ -73,105 +65,93 @@ 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
+NOT_FOUND_SENTINEL = object()
 
 
 def _equal_chance_permutation(
-    objs: Sequence[Item | Album],
+    objs: Sequence[T],
     field: str = "albumartist",
     random_gen: random.Random | None = None,
-) -> Iterable[Item | Album]:
+) -> Iterable[T]:
     """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
+    rand: random.Random = random_gen or random.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 get_attr(obj: T) -> Any:
+        try:
+            return key(obj)
+        except AttributeError:
+            return NOT_FOUND_SENTINEL
+
+    groups: dict[Any, list[T]] = {
+        NOT_FOUND_SENTINEL: [],
+    }
+    for k, values in groupby(objs, key=get_attr):
+        groups[k] = list(values)
+        # shuffle in category
+        rand.shuffle(groups[k])
+
+    # Remove items without the field value.
+    del groups[NOT_FOUND_SENTINEL]
+    while groups:
+        group = rand.choice(list(groups.keys()))
+        yield groups[group].pop()
+        if not groups[group]:
+            del groups[group]
 
 
 def _take_time(
-    iter: Iterable[Item | Album],
+    iter: Iterable[T],
     secs: float,
-) -> list[Item | Album]:
+) -> Iterable[T]:
     """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)
+        length = obj.length
         if total_time + length <= secs:
-            out.append(obj)
+            yield obj
             total_time += length
-    return out
 
 
 def random_objs(
-    objs: Sequence[Item | Album],
-    number=1,
-    time: float | None = None,
+    objs: Sequence[T],
+    number: int = 1,
+    time_minutes: 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.
+) -> Iterable[T]:
+    """Get a random subset of items, optionally constrained by time or count.
+
+    Args:
+    - objs: The sequence of objects to choose from.
+    - number: The number of objects to select.
+    - time_minutes: If specified, the total length of selected objects
+      should not exceed this many minutes.
+    - equal_chance: If True, each artist has the same chance of being
+      selected, regardless of how many tracks they have.
+    - random_gen: An optional random generator to use for shuffling.
     """
-    rand = random_gen or random
+    rand: random.Random = random_gen or random.Random()
 
     # Permute the objects either in a straightforward way or an
     # artist-balanced way.
+    perm: Iterable[T]
     if equal_chance:
-        perm = _equal_chance_permutation(objs)
+        perm = _equal_chance_permutation(objs, random_gen=rand)
     else:
         perm = list(objs)
         rand.shuffle(perm)
 
     # Select objects by time our count.
-    if time:
-        return _take_time(perm, time * 60)
+    if time_minutes:
+        return _take_time(perm, time_minutes * 60)
     else:
-        return _take(perm, number)
+        return islice(perm, number)