random.Random and subclasses: split _randbelow implementation

Splits the getrandbits-dependent and -independent branches of
random.Random._randbelow into separate methods and selects the
implementation to be used by Random and its subclasses at class
creation time for increased performance.

Author: Wolfgang Maier

Lib/random.py

@@ -38,7 +38,7 @@
 """
 
 from warnings import warn as _warn
-from types import MethodType as _MethodType, BuiltinMethodType as _BuiltinMethodType
+from types import FunctionType as _FunctionType
 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 +94,18 @@ def __init__(self, x=None):
         self.seed(x)
         self.gauss_next = None
 
+    def __init_subclass__(cls, **kwargs):
+        # Only call self.getrandbits if the original random() builtin method
+        # has not been overridden or if a new getrandbits() was supplied.
+        if type(cls.__dict__.get('getrandbits')) is _FunctionType:
+            cls._randbelow = cls._randbelow_with_getrandbits
+        elif 'random' in cls.__dict__:
+            # There's an overridden random() method but no new getrandbits() method,
+            # so we can only use random() from here.
+            cls._randbelow = cls._randbelow_without_getrandbits
+        else:
+            cls._randbelow = getattr(cls, cls._randbelow.__name__)
+
     def seed(self, a=None, version=2):
         """Initialize internal state from hashable object.
 
@@ -221,22 +233,26 @@ def randint(self, a, b):
 
         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.
+        """
+
+        if n == 0:
+            raise ValueError("Upper boundary cannot be zero")
+
+        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 +267,8 @@ def _randbelow(self, n, int=int, maxsize=1<<BPF, type=type,
             r = random()
         return int(r*maxsize) % n
 
+    _randbelow = _randbelow_with_getrandbits
+
 ## -------------------- sequence methods  -------------------
 
     def choice(self, seq):

Lib/test/test_random.py

@@ -640,21 +640,22 @@ def test_randbelow_logic(self, _log=log, int=int):
             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 +673,10 @@ def test_randbelow_overridden_random(self, random_mock):
         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 +929,81 @@ def test_betavariate_return_zero(self, gammavariate_mock):
         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_overriding_random(self):
+        # First, let's assert our base class gets the selection of its
+        # _randbelow implementation right.
+        self.assertIs(
+            random.Random._randbelow, random.Random._randbelow_with_getrandbits
+        )
+
+        # Subclasses with a random method that got more recently defined than
+        # their getrandbits method should not rely on getrandbits in
+        # _randbelow, but select the getrandbits-independent implementation
+        # of _randbelow instead.
+
+        # Subclass doesn't override any of the methods => keep using
+        # original getrandbits-dependent version of _randbelow
+        class SubClass1(random.Random):
+            pass
+        self.assertIs(
+            SubClass1._randbelow, random.Random._randbelow_with_getrandbits
+        )
+        # subclass providing its own random **and** getrandbits methods
+        # like random.SystemRandom does => keep relying on getrandbits for
+        # _randbelow
+        class SubClass2(random.Random):
+            def random(self):
+                pass
+
+            def getrandbits(self):
+                pass
+        self.assertIs(
+            SubClass2._randbelow, random.Random._randbelow_with_getrandbits
+        )
+        # subclass providing only random => switch to getrandbits-independent
+        # version of _randbelow
+        class SubClass3(random.Random):
+            def random(self):
+                pass
+        self.assertIs(
+            SubClass3._randbelow, random.Random._randbelow_without_getrandbits
+        )
+        # subclass defining getrandbits to complement its inherited random
+        # => can now rely on getrandbits for _randbelow again
+        class SubClass4(SubClass3):
+            def getrandbits(self):
+                pass
+        self.assertIs(
+            SubClass4._randbelow, random.Random._randbelow_with_getrandbits
+        )
+        # subclass overriding the getrandbits-dependent implementation of
+        # _randbelow => make sure it is used
+        class SubClass5(SubClass4):
+            def _randbelow_with_getrandbits(self):
+                pass
+        self.assertIs(
+            SubClass5._randbelow, SubClass5._randbelow_with_getrandbits
+        )
+        self.assertIsNot(
+            SubClass5._randbelow, random.Random._randbelow_with_getrandbits
+        )
+        # subclass defining random making it more recent than its inherited
+        # getrandbits => switch back to getrandbits-independent implementaion
+        class SubClass6(SubClass5):
+            def random(self):
+                pass
+        self.assertIs(
+            SubClass6._randbelow, random.Random._randbelow_without_getrandbits
+        )
+
 class TestModule(unittest.TestCase):
     def testMagicConstants(self):
         self.assertAlmostEqual(random.NV_MAGICCONST, 1.71552776992141)
@@ -937,13 +1015,6 @@ def test__all__(self):
         # 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