Fixed the randomness of the equal change to have more than one song by artist

This commit is contained in:
Georges Dubus 2013-05-19 17:22:19 +02:00
parent bcae495c50
commit 5dbff4a8eb

View file

@ -21,6 +21,7 @@ from beets.util.functemplate import Template
import random import random
from operator import attrgetter from operator import attrgetter
from itertools import groupby from itertools import groupby
import collections
def random_item(lib, opts, args): def random_item(lib, opts, args):
query = decargs(args) query = decargs(args)
@ -34,18 +35,29 @@ def random_item(lib, opts, args):
objs = list(lib.albums(query=query)) objs = list(lib.albums(query=query))
else: else:
objs = list(lib.items(query=query)) objs = list(lib.items(query=query))
if opts.equal_chance: if opts.equal_chance:
key = attrgetter('albumartist') key = attrgetter('albumartist')
objs.sort(key=key) objs.sort(key=key)
# Now the objs are list of albums or items
objs = [list(v) for k, v in groupby(objs, key)]
number = min(len(objs), opts.number) # {artists: objects}
objs = random.sample(objs, number) objs_by_artists = {artist: list(v) for artist, v in groupby(objs, key)}
artists = objs_by_artists.keys()
if opts.equal_chance: # {artist: count}
# Select one random item from that artist selected_artists = collections.defaultdict(int)
objs = map(random.choice, objs) for _ in range(opts.number):
selected_artists[random.choice(artists)] += 1
objs = []
for artist, count in selected_artists.items():
objs_from_artist = objs_by_artists[artist]
number = min(count, len(objs_from_artist))
objs.extend(random.sample(objs_from_artist, number))
else:
number = min(len(objs), opts.number)
objs = random.sample(objs, number)
for item in objs: for item in objs:
print_obj(item, lib, template) print_obj(item, lib, template)