Skip to content

Commit

Permalink
bpo-42944 Fix Random.sample when counts is not None (pythonGH-24235)
Browse files Browse the repository at this point in the history
  • Loading branch information
jonanifranco committed Jan 18, 2021
1 parent 314b878 commit f7b5bac
Show file tree
Hide file tree
Showing 3 changed files with 29 additions and 28 deletions.
2 changes: 1 addition & 1 deletion Lib/random.py
Expand Up @@ -479,7 +479,7 @@ def sample(self, population, k, *, counts=None):
raise TypeError('Counts must be integers')
if total <= 0:
raise ValueError('Total of counts must be greater than zero')
selections = sample(range(total), k=k)
selections = self.sample(range(total), k=k)
bisect = _bisect
return [population[bisect(cum_counts, s)] for s in selections]
randbelow = self._randbelow
Expand Down
54 changes: 27 additions & 27 deletions Lib/test/test_random.py
Expand Up @@ -223,33 +223,6 @@ def test_sample_with_counts(self):
with self.assertRaises(ValueError):
sample(['red', 'green', 'blue'], counts=[1, 2, 3, 4], k=2) # too many counts

def test_sample_counts_equivalence(self):
# Test the documented strong equivalence to a sample with repeated elements.
# We run this test on random.Random() which makes deterministic selections
# for a given seed value.
sample = random.sample
seed = random.seed

colors = ['red', 'green', 'blue', 'orange', 'black', 'amber']
counts = [500, 200, 20, 10, 5, 1 ]
k = 700
seed(8675309)
s1 = sample(colors, counts=counts, k=k)
seed(8675309)
expanded = [color for (color, count) in zip(colors, counts) for i in range(count)]
self.assertEqual(len(expanded), sum(counts))
s2 = sample(expanded, k=k)
self.assertEqual(s1, s2)

pop = 'abcdefghi'
counts = [10, 9, 8, 7, 6, 5, 4, 3, 2]
seed(8675309)
s1 = ''.join(sample(pop, counts=counts, k=30))
expanded = ''.join([letter for (letter, count) in zip(pop, counts) for i in range(count)])
seed(8675309)
s2 = ''.join(sample(expanded, k=30))
self.assertEqual(s1, s2)

def test_choices(self):
choices = self.gen.choices
data = ['red', 'green', 'blue', 'yellow']
Expand Down Expand Up @@ -957,6 +930,33 @@ def test_randbytes_getrandbits(self):
self.assertEqual(self.gen.randbytes(n),
gen2.getrandbits(n * 8).to_bytes(n, 'little'))

def test_sample_counts_equivalence(self):
# Test the documented strong equivalence to a sample with repeated elements.
# We run this test on random.Random() which makes deterministic selections
# for a given seed value.
sample = self.gen.sample
seed = self.gen.seed

colors = ['red', 'green', 'blue', 'orange', 'black', 'amber']
counts = [500, 200, 20, 10, 5, 1 ]
k = 700
seed(8675309)
s1 = sample(colors, counts=counts, k=k)
seed(8675309)
expanded = [color for (color, count) in zip(colors, counts) for i in range(count)]
self.assertEqual(len(expanded), sum(counts))
s2 = sample(expanded, k=k)
self.assertEqual(s1, s2)

pop = 'abcdefghi'
counts = [10, 9, 8, 7, 6, 5, 4, 3, 2]
seed(8675309)
s1 = ''.join(sample(pop, counts=counts, k=30))
expanded = ''.join([letter for (letter, count) in zip(pop, counts) for i in range(count)])
seed(8675309)
s2 = ''.join(sample(expanded, k=30))
self.assertEqual(s1, s2)


def gamma(z, sqrt2pi=(2.0*pi)**0.5):
# Reflection to right half of complex plane
Expand Down
@@ -0,0 +1 @@
Fix ``random.Random.sample`` when ``counts`` argument is not ``None``.

0 comments on commit f7b5bac

Please sign in to comment.