Skip to content

Commit 8db2fa2

Browse files
[3.12] gh-130285: Fix handling of zero or empty counts in random.sample() (gh-130291) (gh-130417)
1 parent 91e5e24 commit 8db2fa2

File tree

3 files changed

+21
-5
lines changed

3 files changed

+21
-5
lines changed

Lib/random.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -417,11 +417,11 @@ def sample(self, population, k, *, counts=None):
417417
cum_counts = list(_accumulate(counts))
418418
if len(cum_counts) != n:
419419
raise ValueError('The number of counts does not match the population')
420-
total = cum_counts.pop()
420+
total = cum_counts.pop() if cum_counts else 0
421421
if not isinstance(total, int):
422422
raise TypeError('Counts must be integers')
423-
if total <= 0:
424-
raise ValueError('Total of counts must be greater than zero')
423+
if total < 0:
424+
raise ValueError('Counts must be non-negative')
425425
selections = self.sample(range(total), k=k)
426426
bisect = _bisect
427427
return [population[bisect(cum_counts, s)] for s in selections]

Lib/test/test_random.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -224,15 +224,27 @@ def test_sample_with_counts(self):
224224
sample(['red', 'green', 'blue'], counts=10, k=10) # counts not iterable
225225
with self.assertRaises(ValueError):
226226
sample(['red', 'green', 'blue'], counts=[-3, -7, -8], k=2) # counts are negative
227-
with self.assertRaises(ValueError):
228-
sample(['red', 'green', 'blue'], counts=[0, 0, 0], k=2) # counts are zero
229227
with self.assertRaises(ValueError):
230228
sample(['red', 'green'], counts=[10, 10], k=21) # population too small
231229
with self.assertRaises(ValueError):
232230
sample(['red', 'green', 'blue'], counts=[1, 2], k=2) # too few counts
233231
with self.assertRaises(ValueError):
234232
sample(['red', 'green', 'blue'], counts=[1, 2, 3, 4], k=2) # too many counts
235233

234+
# Cases with zero counts match equivalents without counts (see gh-130285)
235+
self.assertEqual(
236+
sample('abc', k=0, counts=[0, 0, 0]),
237+
sample([], k=0),
238+
)
239+
self.assertEqual(
240+
sample([], 0, counts=[]),
241+
sample([], 0),
242+
)
243+
with self.assertRaises(ValueError):
244+
sample([], 1, counts=[])
245+
with self.assertRaises(ValueError):
246+
sample('x', 1, counts=[0])
247+
236248
def test_choices(self):
237249
choices = self.gen.choices
238250
data = ['red', 'green', 'blue', 'yellow']
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
Fix corner case for :func:`random.sample` allowing the *counts* parameter to
2+
specify an empty population. So now, ``sample([], 0, counts=[])`` and
3+
``sample('abc', k=0, counts=[0, 0, 0])`` both give the same result as
4+
``sample([], 0)``.

0 commit comments

Comments
 (0)