diff options
Diffstat (limited to 'Lib')
| -rw-r--r-- | Lib/random.py | 17 | ||||
| -rw-r--r-- | Lib/test/test_random.py | 20 |
2 files changed, 32 insertions, 5 deletions
diff --git a/Lib/random.py b/Lib/random.py index 943fa51..465f477 100644 --- a/Lib/random.py +++ b/Lib/random.py @@ -285,6 +285,15 @@ class Random(_random.Random): large population: sample(xrange(10000000), 60) """ + # XXX Although the documentation says `population` is "a sequence", + # XXX attempts are made to cater to any iterable with a __len__ + # XXX method. This has had mixed success. Examples from both + # XXX sides: sets work fine, and should become officially supported; + # XXX dicts are much harder, and have failed in various subtle + # XXX ways across attempts. Support for mapping types should probably + # XXX be dropped (and users should pass mapping.keys() or .values() + # XXX explicitly). + # Sampling without replacement entails tracking either potential # selections (the pool) in a list or previous selections in a set. @@ -304,7 +313,9 @@ class Random(_random.Random): setsize = 21 # size of a small set minus size of an empty list if k > 5: setsize += 4 ** _ceil(_log(k * 3, 4)) # table size for big sets - if n <= setsize: # is an n-length list smaller than a k-length set + if n <= setsize or hasattr(population, "keys"): + # An n-length list is smaller than a k-length set, or this is a + # mapping type so the other algorithm wouldn't work. pool = list(population) for i in xrange(k): # invariant: non-selected at [0,n-i) j = _int(random() * (n-i)) @@ -320,10 +331,10 @@ class Random(_random.Random): j = _int(random() * n) selected_add(j) result[i] = population[j] - except (TypeError, KeyError): # handle sets and dictionaries + except (TypeError, KeyError): # handle (at least) sets if isinstance(population, list): raise - return self.sample(list(population), k) + return self.sample(tuple(population), k) return result ## -------------------- real-valued distributions ------------------- diff --git a/Lib/test/test_random.py b/Lib/test/test_random.py index c9431b3..bba4c7c 100644 --- a/Lib/test/test_random.py +++ b/Lib/test/test_random.py @@ -93,12 +93,28 @@ class TestBasicOps(unittest.TestCase): self.gen.sample(set(range(20)), 2) self.gen.sample(range(20), 2) self.gen.sample(xrange(20), 2) - self.gen.sample(dict.fromkeys('abcdefghijklmnopqrst'), 2) self.gen.sample(str('abcdefghijklmnopqrst'), 2) self.gen.sample(tuple('abcdefghijklmnopqrst'), 2) + + def test_sample_on_dicts(self): + self.gen.sample(dict.fromkeys('abcdefghijklmnopqrst'), 2) + # SF bug #1460340 -- random.sample can raise KeyError a = dict.fromkeys(range(10)+range(10,100,2)+range(100,110)) - self.gen.sample(a,3) + self.gen.sample(a, 3) + + # A followup to bug #1460340: sampling from a dict could return + # a subset of its keys or of its values, depending on the size of + # the subset requested. + N = 30 + d = dict((i, complex(i, i)) for i in xrange(N)) + for k in xrange(N+1): + samp = self.gen.sample(d, k) + # Verify that we got ints back (keys); the values are complex. + for x in samp: + self.assert_(type(x) is int) + samp.sort() + self.assertEqual(samp, range(N)) def test_gauss(self): # Ensure that the seed() method initializes all the hidden state. In |