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

1 files changed, 38 insertions, 14 deletions

diff --git a/Lib/random.py b/Lib/random.py
index 0bc24174e13..0ed5511e9f6 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):