BUG/ENH: ALlow bit generators to supply their own ctor

Allow bit generators to supply their own constructors to enable Generator
objects using arbitrary bit generators to be supported

closes #22012

1 files changed, 23 insertions, 26 deletions

diff --git a/numpy/random/_pickle.py b/numpy/random/_pickle.py
index 5e89071e8..073993726 100644
--- a/numpy/random/_pickle.py
+++ b/numpy/random/_pickle.py
@@ -14,19 +14,19 @@ BitGenerators = {'MT19937': MT19937,
                  }
 
 
-def __generator_ctor(bit_generator_name='MT19937'):
+def __bit_generator_ctor(bit_generator_name='MT19937'):
     """
-    Pickling helper function that returns a Generator object
+    Pickling helper function that returns a bit generator object
 
     Parameters
     ----------
     bit_generator_name : str
-        String containing the core BitGenerator
+        String containing the name of the BitGenerator
 
     Returns
     -------
-    rg : Generator
-        Generator using the named core BitGenerator
+    bit_generator : BitGenerator
+        BitGenerator instance
     """
     if bit_generator_name in BitGenerators:
         bit_generator = BitGenerators[bit_generator_name]
@@ -34,50 +34,47 @@ def __generator_ctor(bit_generator_name='MT19937'):
         raise ValueError(str(bit_generator_name) + ' is not a known '
                                                    'BitGenerator module.')
 
-    return Generator(bit_generator())
+    return bit_generator()
 
 
-def __bit_generator_ctor(bit_generator_name='MT19937'):
+def __generator_ctor(bit_generator_name="MT19937",
+                     bit_generator_ctor=__bit_generator_ctor):
     """
-    Pickling helper function that returns a bit generator object
+    Pickling helper function that returns a Generator object
 
     Parameters
     ----------
     bit_generator_name : str
-        String containing the name of the BitGenerator
+        String containing the core BitGenerator's name
+    bit_generator_ctor : callable, optional
+        Callable function that takes bit_generator_name as its only argument
+        and returns an instantized bit generator.
 
     Returns
     -------
-    bit_generator : BitGenerator
-        BitGenerator instance
+    rg : Generator
+        Generator using the named core BitGenerator
     """
-    if bit_generator_name in BitGenerators:
-        bit_generator = BitGenerators[bit_generator_name]
-    else:
-        raise ValueError(str(bit_generator_name) + ' is not a known '
-                                                   'BitGenerator module.')
-
-    return bit_generator()
+    return Generator(bit_generator_ctor(bit_generator_name))
 
 
-def __randomstate_ctor(bit_generator_name='MT19937'):
+def __randomstate_ctor(bit_generator_name="MT19937",
+                       bit_generator_ctor=__bit_generator_ctor):
     """
     Pickling helper function that returns a legacy RandomState-like object
 
     Parameters
     ----------
     bit_generator_name : str
-        String containing the core BitGenerator
+        String containing the core BitGenerator's name
+    bit_generator_ctor : callable, optional
+        Callable function that takes bit_generator_name as its only argument
+        and returns an instantized bit generator.
 
     Returns
     -------
     rs : RandomState
         Legacy RandomState using the named core BitGenerator
     """
-    if bit_generator_name in BitGenerators:
-        bit_generator = BitGenerators[bit_generator_name]
-    else:
-        raise ValueError(str(bit_generator_name) + ' is not a known '
-                                                   'BitGenerator module.')
 
-    return RandomState(bit_generator())
+    return RandomState(bit_generator_ctor(bit_generator_name))