Skip to content

Commit

Permalink
pythongh-130285: Fix handling of zero or empty counts in random.sampl…
Browse files Browse the repository at this point in the history
…e() (pythongh-130291)

(cherry picked from commit 286c517)

Co-authored-by: Raymond Hettinger <rhettinger@users.noreply.github.com>
  • Loading branch information
rhettinger authored and miss-islington committed Feb 21, 2025
1 parent 91e5e24 commit 4e0e702
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 5 deletions.
6 changes: 3 additions & 3 deletions Lib/random.py
Original file line number Diff line number Diff line change
Expand Up @@ -417,11 +417,11 @@ def sample(self, population, k, *, counts=None):
cum_counts = list(_accumulate(counts))
if len(cum_counts) != n:
raise ValueError('The number of counts does not match the population')
total = cum_counts.pop()
total = cum_counts.pop() if cum_counts else 0
if not isinstance(total, int):
raise TypeError('Counts must be integers')
if total <= 0:
raise ValueError('Total of counts must be greater than zero')
if total < 0:
raise ValueError('Counts must be non-negative')
selections = self.sample(range(total), k=k)
bisect = _bisect
return [population[bisect(cum_counts, s)] for s in selections]
Expand Down
16 changes: 14 additions & 2 deletions Lib/test/test_random.py
Original file line number Diff line number Diff line change
Expand Up @@ -224,15 +224,27 @@ def test_sample_with_counts(self):
sample(['red', 'green', 'blue'], counts=10, k=10) # counts not iterable
with self.assertRaises(ValueError):
sample(['red', 'green', 'blue'], counts=[-3, -7, -8], k=2) # counts are negative
with self.assertRaises(ValueError):
sample(['red', 'green', 'blue'], counts=[0, 0, 0], k=2) # counts are zero
with self.assertRaises(ValueError):
sample(['red', 'green'], counts=[10, 10], k=21) # population too small
with self.assertRaises(ValueError):
sample(['red', 'green', 'blue'], counts=[1, 2], k=2) # too few counts
with self.assertRaises(ValueError):
sample(['red', 'green', 'blue'], counts=[1, 2, 3, 4], k=2) # too many counts

# Cases with zero counts match equivalents without counts (see gh-130285)
self.assertEqual(
sample('abc', k=0, counts=[0, 0, 0]),
sample([], k=0),
)
self.assertEqual(
sample([], 0, counts=[]),
sample([], 0),
)
with self.assertRaises(ValueError):
sample([], 1, counts=[])
with self.assertRaises(ValueError):
sample('x', 1, counts=[0])

def test_choices(self):
choices = self.gen.choices
data = ['red', 'green', 'blue', 'yellow']
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
Fix corner case for :func:`random.sample` allowing the *counts* parameter to
specify an empty population. So now, ``sample([], 0, counts=[])`` and
``sample('abc', k=0, counts=[0, 0, 0])`` both give the same result as
``sample([], 0)``.

0 comments on commit 4e0e702

Please # to comment.