diff options
Diffstat (limited to 'Lib/test/test_random.py')
| -rw-r--r-- | Lib/test/test_random.py | 112 |
1 files changed, 112 insertions, 0 deletions
diff --git a/Lib/test/test_random.py b/Lib/test/test_random.py index e80ed17a8c..5b6a4f06ba 100644 --- a/Lib/test/test_random.py +++ b/Lib/test/test_random.py @@ -7,6 +7,7 @@ import warnings from functools import partial from math import log, exp, pi, fsum, sin from test import support +from fractions import Fraction class TestBasicOps: # Superclass with tests common to all generators. @@ -109,6 +110,7 @@ class TestBasicOps: self.assertEqual(self.gen.sample([], 0), []) # test edge case N==k==0 # Exception raised if size of sample exceeds that of population self.assertRaises(ValueError, self.gen.sample, population, N+1) + self.assertRaises(ValueError, self.gen.sample, [], -1) def test_sample_distribution(self): # For the entire allowable range of 0 <= k <= N, validate that @@ -141,6 +143,83 @@ class TestBasicOps: def test_sample_on_dicts(self): self.assertRaises(TypeError, self.gen.sample, dict.fromkeys('abcdef'), 2) + def test_choices(self): + choices = self.gen.choices + data = ['red', 'green', 'blue', 'yellow'] + str_data = 'abcd' + range_data = range(4) + set_data = set(range(4)) + + # basic functionality + for sample in [ + choices(data, k=5), + choices(data, range(4), k=5), + choices(k=5, population=data, weights=range(4)), + choices(k=5, population=data, cum_weights=range(4)), + ]: + self.assertEqual(len(sample), 5) + self.assertEqual(type(sample), list) + self.assertTrue(set(sample) <= set(data)) + + # test argument handling + with self.assertRaises(TypeError): # missing arguments + choices(2) + + self.assertEqual(choices(data, k=0), []) # k == 0 + self.assertEqual(choices(data, k=-1), []) # negative k behaves like ``[0] * -1`` + with self.assertRaises(TypeError): + choices(data, k=2.5) # k is a float + + self.assertTrue(set(choices(str_data, k=5)) <= set(str_data)) # population is a string sequence + self.assertTrue(set(choices(range_data, k=5)) <= set(range_data)) # population is a range + with self.assertRaises(TypeError): + choices(set_data, k=2) # population is not a sequence + + self.assertTrue(set(choices(data, None, k=5)) <= set(data)) # weights is None + self.assertTrue(set(choices(data, weights=None, k=5)) <= set(data)) + with self.assertRaises(ValueError): + choices(data, [1,2], k=5) # len(weights) != len(population) + with self.assertRaises(TypeError): + choices(data, 10, k=5) # non-iterable weights + with self.assertRaises(TypeError): + choices(data, [None]*4, k=5) # non-numeric weights + for weights in [ + [15, 10, 25, 30], # integer weights + [15.1, 10.2, 25.2, 30.3], # float weights + [Fraction(1, 3), Fraction(2, 6), Fraction(3, 6), Fraction(4, 6)], # fractional weights + [True, False, True, False] # booleans (include / exclude) + ]: + self.assertTrue(set(choices(data, weights, k=5)) <= set(data)) + + with self.assertRaises(ValueError): + choices(data, cum_weights=[1,2], k=5) # len(weights) != len(population) + with self.assertRaises(TypeError): + choices(data, cum_weights=10, k=5) # non-iterable cum_weights + with self.assertRaises(TypeError): + choices(data, cum_weights=[None]*4, k=5) # non-numeric cum_weights + with self.assertRaises(TypeError): + choices(data, range(4), cum_weights=range(4), k=5) # both weights and cum_weights + for weights in [ + [15, 10, 25, 30], # integer cum_weights + [15.1, 10.2, 25.2, 30.3], # float cum_weights + [Fraction(1, 3), Fraction(2, 6), Fraction(3, 6), Fraction(4, 6)], # fractional cum_weights + ]: + self.assertTrue(set(choices(data, cum_weights=weights, k=5)) <= set(data)) + + # Test weight focused on a single element of the population + self.assertEqual(choices('abcd', [1, 0, 0, 0]), ['a']) + self.assertEqual(choices('abcd', [0, 1, 0, 0]), ['b']) + self.assertEqual(choices('abcd', [0, 0, 1, 0]), ['c']) + self.assertEqual(choices('abcd', [0, 0, 0, 1]), ['d']) + + # Test consistency with random.choice() for empty population + with self.assertRaises(IndexError): + choices([], k=1) + with self.assertRaises(IndexError): + choices([], weights=[], k=1) + with self.assertRaises(IndexError): + choices([], cum_weights=[], k=5) + def test_gauss(self): # Ensure that the seed() method initializes all the hidden state. In # particular, through 2.2.1 it failed to reset a piece of state used @@ -551,6 +630,39 @@ class MersenneTwister_TestBasicOps(TestBasicOps, unittest.TestCase): self.assertTrue(stop < x <= start) self.assertEqual((x+stop)%step, 0) + def test_choices_algorithms(self): + # The various ways of specifying weights should produce the same results + choices = self.gen.choices + n = 104729 + + self.gen.seed(8675309) + a = self.gen.choices(range(n), k=10000) + + self.gen.seed(8675309) + b = self.gen.choices(range(n), [1]*n, k=10000) + self.assertEqual(a, b) + + self.gen.seed(8675309) + c = self.gen.choices(range(n), cum_weights=range(1, n+1), k=10000) + self.assertEqual(a, c) + + # Amerian Roulette + population = ['Red', 'Black', 'Green'] + weights = [18, 18, 2] + cum_weights = [18, 36, 38] + expanded_population = ['Red'] * 18 + ['Black'] * 18 + ['Green'] * 2 + + self.gen.seed(9035768) + a = self.gen.choices(expanded_population, k=10000) + + self.gen.seed(9035768) + b = self.gen.choices(population, weights, k=10000) + self.assertEqual(a, b) + + self.gen.seed(9035768) + c = self.gen.choices(population, cum_weights=cum_weights, k=10000) + self.assertEqual(a, c) + def gamma(z, sqrt2pi=(2.0*pi)**0.5): # Reflection to right half of complex plane if z < 0.5: |