Adding unlimited=bool to files_varied -- if set to true, it will never issue StopIteration

This commit is contained in:
Wesley Spikes 2011-07-18 15:26:03 -07:00
parent 43b2f8695d
commit 44cb2f256e

View file

@ -160,7 +160,7 @@ def names(mean, stddev, charset=None, seed=None):
name = ''.join(rand.choice(charset) for _ in xrange(length))
yield name
def files_varied(groups):
def files_varied(groups, unlimited=False):
""" Yields a weighted-random selection of file-like objects. """
# Quick data type sanity.
assert groups and isinstance(groups, (list, tuple))
@ -193,12 +193,13 @@ def files_varied(groups):
num -= file_set.num
continue
total_num -= 1
file_set.num -= 1
if not unlimited:
total_num -= 1
file_set.num -= 1
# None left in this set!
if file_set.num == 0:
file_sets.remove(file_set)
# None left in this set!
if file_set.num == 0:
file_sets.remove(file_set)
ok = 1
yield next(file_set.files)