bpo-33144: random.Random and subclasses: split _randbelow implementation (GH-6291)

3 files changed, 107 insertions, 29 deletions

diff --git a/Lib/random.py b/Lib/random.py
index 0bc2417..0ed5511 100644
--- a/Lib/random.py
+++ b/Lib/random.py
@@ -38,7 +38,6 @@ General notes on the underlying Mersenne Twister core generator:
 """
 
 from warnings import warn as _warn
-from types import MethodType as _MethodType, BuiltinMethodType as _BuiltinMethodType
 from math import log as _log, exp as _exp, pi as _pi, e as _e, ceil as _ceil
 from math import sqrt as _sqrt, acos as _acos, cos as _cos, sin as _sin
 from os import urandom as _urandom
@@ -94,6 +93,28 @@ class Random(_random.Random):
         self.seed(x)
         self.gauss_next = None
 
+    def __init_subclass__(cls, **kwargs):
+        """Control how subclasses generate random integers.
+
+        The algorithm a subclass can use depends on the random() and/or
+        getrandbits() implementation available to it and determines
+        whether it can generate random integers from arbitrarily large
+        ranges.
+        """
+
+        if (cls.random is _random.Random.random) or (
+            cls.getrandbits is not _random.Random.getrandbits):
+            # The original random() builtin method has not been overridden
+            # or a new getrandbits() was supplied.
+            # The subclass can use the getrandbits-dependent implementation
+            # of _randbelow().
+            cls._randbelow = cls._randbelow_with_getrandbits
+        else:
+            # There's an overridden random() method but no new getrandbits(),
+            # so the subclass can only use the getrandbits-independent
+            # implementation of _randbelow().
+            cls._randbelow = cls._randbelow_without_getrandbits
+
     def seed(self, a=None, version=2):
         """Initialize internal state from hashable object.
 
@@ -221,22 +242,23 @@ class Random(_random.Random):
 
         return self.randrange(a, b+1)
 
-    def _randbelow(self, n, int=int, maxsize=1<<BPF, type=type,
-                   Method=_MethodType, BuiltinMethod=_BuiltinMethodType):
+    def _randbelow_with_getrandbits(self, n):
         "Return a random int in the range [0,n).  Raises ValueError if n==0."
 
-        random = self.random
         getrandbits = self.getrandbits
-        # Only call self.getrandbits if the original random() builtin method
-        # has not been overridden or if a new getrandbits() was supplied.
-        if type(random) is BuiltinMethod or type(getrandbits) is Method:
-            k = n.bit_length()  # don't use (n-1) here because n can be 1
-            r = getrandbits(k)          # 0 <= r < 2**k
-            while r >= n:
-                r = getrandbits(k)
-            return r
-        # There's an overridden random() method but no new getrandbits() method,
-        # so we can only use random() from here.
+        k = n.bit_length()  # don't use (n-1) here because n can be 1
+        r = getrandbits(k)          # 0 <= r < 2**k
+        while r >= n:
+            r = getrandbits(k)
+        return r
+
+    def _randbelow_without_getrandbits(self, n, int=int, maxsize=1<<BPF):
+        """Return a random int in the range [0,n).  Raises ValueError if n==0.
+
+        The implementation does not use getrandbits, but only random.
+        """
+
+        random = self.random
         if n >= maxsize:
             _warn("Underlying random() generator does not supply \n"
                 "enough bits to choose from a population range this large.\n"
@@ -251,6 +273,8 @@ class Random(_random.Random):
             r = random()
         return int(r*maxsize) % n
 
+    _randbelow = _randbelow_with_getrandbits
+
 ## -------------------- sequence methods  -------------------
 
     def choice(self, seq):
diff --git a/Lib/test/test_random.py b/Lib/test/test_random.py
index eee245d..d91908b 100644
--- a/Lib/test/test_random.py
+++ b/Lib/test/test_random.py
@@ -5,6 +5,7 @@ import os
 import time
 import pickle
 import warnings
+import logging
 from functools import partial
 from math import log, exp, pi, fsum, sin, factorial
 from test import support
@@ -619,6 +620,16 @@ class MersenneTwister_TestBasicOps(TestBasicOps, unittest.TestCase):
         self.assertRaises(ValueError, self.gen.getrandbits, 0)
         self.assertRaises(ValueError, self.gen.getrandbits, -1)
 
+    def test_randrange_uses_getrandbits(self):
+        # Verify use of getrandbits by randrange
+        # Use same seed as in the cross-platform repeatability test
+        # in test_genrandbits above.
+        self.gen.seed(1234567)
+        # If randrange uses getrandbits, it should pick getrandbits(100)
+        # when called with a 100-bits stop argument.
+        self.assertEqual(self.gen.randrange(2**99),
+                         97904845777343510404718956115)
+
     def test_randbelow_logic(self, _log=log, int=int):
         # check bitcount transition points:  2**i and 2**(i+1)-1
         # show that: k = int(1.001 + _log(n, 2))
@@ -640,21 +651,22 @@ class MersenneTwister_TestBasicOps(TestBasicOps, unittest.TestCase):
             self.assertEqual(k, numbits)        # note the stronger assertion
             self.assertTrue(2**k > n > 2**(k-1))   # note the stronger assertion
 
-    @unittest.mock.patch('random.Random.random')
-    def test_randbelow_overridden_random(self, random_mock):
+    def test_randbelow_without_getrandbits(self):
         # Random._randbelow() can only use random() when the built-in one
         # has been overridden but no new getrandbits() method was supplied.
-        random_mock.side_effect = random.SystemRandom().random
         maxsize = 1<<random.BPF
         with warnings.catch_warnings():
             warnings.simplefilter("ignore", UserWarning)
             # Population range too large (n >= maxsize)
-            self.gen._randbelow(maxsize+1, maxsize = maxsize)
-        self.gen._randbelow(5640, maxsize = maxsize)
+            self.gen._randbelow_without_getrandbits(
+                maxsize+1, maxsize=maxsize
+            )
+        self.gen._randbelow_without_getrandbits(5640, maxsize=maxsize)
         # issue 33203: test that _randbelow raises ValueError on
         # n == 0 also in its getrandbits-independent branch.
         with self.assertRaises(ValueError):
-            self.gen._randbelow(0, maxsize=maxsize)
+            self.gen._randbelow_without_getrandbits(0, maxsize=maxsize)
+
         # This might be going too far to test a single line, but because of our
         # noble aim of achieving 100% test coverage we need to write a case in
         # which the following line in Random._randbelow() gets executed:
@@ -672,8 +684,10 @@ class MersenneTwister_TestBasicOps(TestBasicOps, unittest.TestCase):
         n = 42
         epsilon = 0.01
         limit = (maxsize - (maxsize % n)) / maxsize
-        random_mock.side_effect = [limit + epsilon, limit - epsilon]
-        self.gen._randbelow(n, maxsize = maxsize)
+        with unittest.mock.patch.object(random.Random, 'random') as random_mock:
+            random_mock.side_effect = [limit + epsilon, limit - epsilon]
+            self.gen._randbelow_without_getrandbits(n, maxsize=maxsize)
+            self.assertEqual(random_mock.call_count, 2)
 
     def test_randrange_bug_1590891(self):
         start = 1000000000000
@@ -926,6 +940,49 @@ class TestDistributions(unittest.TestCase):
         gammavariate_mock.return_value = 0.0
         self.assertEqual(0.0, random.betavariate(2.71828, 3.14159))
 
+class TestRandomSubclassing(unittest.TestCase):
+    def test_random_subclass_with_kwargs(self):
+        # SF bug #1486663 -- this used to erroneously raise a TypeError
+        class Subclass(random.Random):
+            def __init__(self, newarg=None):
+                random.Random.__init__(self)
+        Subclass(newarg=1)
+
+    def test_subclasses_overriding_methods(self):
+        # Subclasses with an overridden random, but only the original
+        # getrandbits method should not rely on getrandbits in for randrange,
+        # but should use a getrandbits-independent implementation instead.
+
+        # subclass providing its own random **and** getrandbits methods
+        # like random.SystemRandom does => keep relying on getrandbits for
+        # randrange
+        class SubClass1(random.Random):
+            def random(self):
+                return super().random()
+
+            def getrandbits(self, n):
+                logging.getLogger('getrandbits').info('used getrandbits')
+                return super().getrandbits(n)
+        with self.assertLogs('getrandbits'):
+            SubClass1().randrange(42)
+
+        # subclass providing only random => can only use random for randrange
+        class SubClass2(random.Random):
+            def random(self):
+                logging.getLogger('random').info('used random')
+                return super().random()
+        with self.assertLogs('random'):
+            SubClass2().randrange(42)
+
+        # subclass defining getrandbits to complement its inherited random
+        # => can now rely on getrandbits for randrange again
+        class SubClass3(SubClass2):
+            def getrandbits(self, n):
+                logging.getLogger('getrandbits').info('used getrandbits')
+                return super().getrandbits(n)
+        with self.assertLogs('getrandbits'):
+            SubClass3().randrange(42)
+
 class TestModule(unittest.TestCase):
     def testMagicConstants(self):
         self.assertAlmostEqual(random.NV_MAGICCONST, 1.71552776992141)
@@ -937,13 +994,6 @@ class TestModule(unittest.TestCase):
         # tests validity but not completeness of the __all__ list
         self.assertTrue(set(random.__all__) <= set(dir(random)))
 
-    def test_random_subclass_with_kwargs(self):
-        # SF bug #1486663 -- this used to erroneously raise a TypeError
-        class Subclass(random.Random):
-            def __init__(self, newarg=None):
-                random.Random.__init__(self)
-        Subclass(newarg=1)
-
     @unittest.skipUnless(hasattr(os, "fork"), "fork() required")
     def test_after_fork(self):
         # Test the global Random instance gets reseeded in child
diff --git a/Misc/NEWS.d/next/Library/2018-04-10-14-50-30.bpo-33144.iZr4et.rst b/Misc/NEWS.d/next/Library/2018-04-10-14-50-30.bpo-33144.iZr4et.rst
new file mode 100644
index 0000000..eb6b9b7f
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2018-04-10-14-50-30.bpo-33144.iZr4et.rst
@@ -0,0 +1,4 @@
+``random.Random()`` and its subclassing mechanism got optimized to check only
+once at class/subclass instantiation time whether its ``getrandbits()`` method
+can be relied on by other methods, including ``randrange()``, for the
+generation of arbitrarily large random integers.  Patch by Wolfgang Maier.