Repository: numpy/numpy
Pull request: BUG/ENH: Allow bit generators to supply their own constructor
Allow bit generators to supply their own constructors to enable Generator
objects using arbitrary bit generators to be supported

closes #22012

<!--         ----------------------------------------------------------------
                MAKE SURE YOUR PR GETS THE ATTENTION IT DESERVES!
                ----------------------------------------------------------------

*  FORMAT IT RIGHT:
      https://www.numpy.org/devdocs/dev/development_workflow.html#writing-the-commit-message

*  IF IT'S A NEW FEATURE OR API CHANGE, TEST THE WATERS:
      https://www.numpy.org/devdocs/dev/development_workflow.html#get-the-mailing-list-s-opinion

*  HIT ALL THE GUIDELINES:
      https://numpy.org/devdocs/dev/index.html#guidelines

*  WHAT TO DO IF WE HAVEN'T GOTTEN BACK TO YOU:
      https://www.numpy.org/devdocs/dev/development_workflow.html#getting-your-pr-reviewed
-->


Problem statement:
BUG: np.random.Generator cannot be unpickled using 3rd party bit generators
### Describe the issue:

The unpickling mechanism uses a hard-coded list of bit generators when unpickling a generator instance.  Because of this, it cannot  be used with 3rd party bit generators, or, if the name of the 3rd party bit generator is the same as the NumPy bit generator, it will unpickle using a different class than the original bit generator class

### Reproduce the code example:

```python
import numpy as np
import randomgen as rg
import pickle

# Generator(MT64)
g = np.random.Generator(rg.MT64())
pickle.loads(pickle.dumps(g))
```


### Error message:

```shell
Traceback (most recent call last):
  File "C:\git\randomgen\exp.py", line 7, in <module>
    pickle.loads(pickle.dumps(g))
  File "c:\git\numpy\numpy\random\_pickle.py", line 52, in __generator_ctor
    raise ValueError(
ValueError: MT64  is not a known NumPy BitGenerator class.
```


### NumPy/Python version information:

1.24.0.dev0+559.g5ba36b700 3.10.4 | packaged by conda-forge | (main, Mar 30 2022, 08:38:02) [MSC v.1916 64 bit (AMD64)]

Hints from issue discussion:
# Possible Solutions

## Let 3rd parts register bit generators
Provide an official mechanism to add bit generators to the lookup table (a dict) so that the current infrastructure could be extended.

### Advantages
* Continues to use same mechanism that has been around forever

### Disadvantages
* Could have name conflicts with existing modules in look up table, unless protected
* Requires explicit action by end users or library writers on module import
* Adds to API


 Something like

```python
np.random.register_bit_generator(bitgen: BitGenerator, name: Optional[str] = None):
   """
    Register a bit generator for pickling/unpickling

    Parameters
    ----------
    bitgen : Bit Generator
        The bit generator to register
    name : str, optional
        The name to use in the look up table.  If not provided, bitgen.__name__ is used
    """
```

## Support looking up fully qualified names
Let 3rd party bit generators name them selves using the pattern `module.class` so that the constructor can use `importlib`

### Advantages
* Totally transparent
* No API space
* Name guaranteed to be unique

### Disadvantages
* Hard to test
* Requires 3rd parties to use a specific pattern (pretty minor IMO)
* Requires a FQN, so unlikely to work with locally defined classes (but usually a problem for pickle)

Something like

```python
def __generator_ctor(bit_generator_name:str='MT19937') -> Generator:
    """
    Pickling helper function that returns a Generator object

    Parameters
    ----------
    bit_generator_name : str
        String containing the core BitGenerator. Can also be the
        full path (module.class) of a third-party bit generator.

    Returns
    -------
    rg : Generator
        Generator using the named core BitGenerator
    """
    bit_generator: Optional[Type[BitGenerator]] = None
    if bit_generator_name in BitGenerators:
        bit_generator = BitGenerators[bit_generator_name]
    elif "." in bit_generator_name:
        parts = bit_generator_name.split(".")
        mod_name = ".".join(parts[:-1])
        bit_generator_name = parts[-1]

        import importlib
        try:
            mod = importlib.import_module(mod_name)
            bit_generator = getattr(mod, bit_generator_name, None)
        except ImportError:
            pass

    if bit_generator is None:
        raise ValueError(
            f"{bit_generator_name}  is not a known NumPy BitGenerator class "
            f"and cannot be looked up using it's fully qualified name."
        )
    assert bit_generator is not None

    return Generator(bit_generator())
```

Commit message: 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

Patch to review:
diff --git a/numpy/random/_generator.pyx b/numpy/random/_generator.pyx
index 77f70e9fec7a..da8ab64e2217 100644
--- a/numpy/random/_generator.pyx
+++ b/numpy/random/_generator.pyx
@@ -220,8 +220,11 @@ cdef class Generator:
         self.bit_generator.state = state
 
     def __reduce__(self):
+        ctor, name_tpl, state = self._bit_generator.__reduce__()
+
         from ._pickle import __generator_ctor
-        return __generator_ctor, (self.bit_generator.state['bit_generator'],), self.bit_generator.state
+        # Requirements of __generator_ctor are (name, ctor)
+        return __generator_ctor, (name_tpl[0], ctor), state
 
     @property
     def bit_generator(self):
diff --git a/numpy/random/_pickle.py b/numpy/random/_pickle.py
index 5e89071e88ab..073993726eb3 100644
--- a/numpy/random/_pickle.py
+++ b/numpy/random/_pickle.py
@@ -14,19 +14,19 @@
                  }
 
 
-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))
diff --git a/numpy/random/mtrand.pyx b/numpy/random/mtrand.pyx
index 19d23f6a856b..c9cdb5839a04 100644
--- a/numpy/random/mtrand.pyx
+++ b/numpy/random/mtrand.pyx
@@ -213,9 +213,10 @@ cdef class RandomState:
         self.set_state(state)
 
     def __reduce__(self):
-        state = self.get_state(legacy=False)
+        ctor, name_tpl, _ = self._bit_generator.__reduce__()
+
         from ._pickle import __randomstate_ctor
-        return __randomstate_ctor, (state['bit_generator'],), state
+        return __randomstate_ctor, (name_tpl[0], ctor), self.get_state(legacy=False)
 
     cdef _reset_gauss(self):
         self._aug_state.has_gauss = 0
diff --git a/numpy/random/tests/test_generator_mt19937.py b/numpy/random/tests/test_generator_mt19937.py
index fa55ac0ee96a..0f4437bcba26 100644
--- a/numpy/random/tests/test_generator_mt19937.py
+++ b/numpy/random/tests/test_generator_mt19937.py
@@ -2695,3 +2695,16 @@ def test_contig_req_out(dist, order, dtype):
     assert variates is out
     variates = dist(out=out, dtype=dtype, size=out.shape)
     assert variates is out
+
+
+def test_generator_ctor_old_style_pickle():
+    rg = np.random.Generator(np.random.PCG64DXSM(0))
+    rg.standard_normal(1)
+    # Directly call reduce which is used in pickline
+    ctor, args, state_a = rg.__reduce__()
+    # Simulate unpickling an old pickle that only has the name
+    assert args[:1] == ("PCG64DXSM",)
+    b = ctor(*args[:1])
+    b.bit_generator.state = state_a
+    state_b = b.bit_generator.state
+    assert state_a == state_b
diff --git a/numpy/random/tests/test_randomstate.py b/numpy/random/tests/test_randomstate.py
index 861813a95d1f..aadfe3d78c8f 100644
--- a/numpy/random/tests/test_randomstate.py
+++ b/numpy/random/tests/test_randomstate.py
@@ -2020,3 +2020,21 @@ def test_broadcast_size_error():
         random.binomial([1, 2], 0.3, size=(2, 1))
     with pytest.raises(ValueError):
         random.binomial([1, 2], [0.3, 0.7], size=(2, 1))
+
+
+def test_randomstate_ctor_old_style_pickle():
+    rs = np.random.RandomState(MT19937(0))
+    rs.standard_normal(1)
+    # Directly call reduce which is used in pickline
+    ctor, args, state_a = rs.__reduce__()
+    # Simulate unpickling an old pickle that only has the name
+    assert args[:1] == ("MT19937",)
+    b = ctor(*args[:1])
+    b.set_state(state_a)
+    state_b = b.get_state(legacy=False)
+
+    assert_equal(state_a['bit_generator'], state_b['bit_generator'])
+    assert_array_equal(state_a['state']['key'], state_b['state']['key'])
+    assert_array_equal(state_a['state']['pos'], state_b['state']['pos'])
+    assert_equal(state_a['has_gauss'], state_b['has_gauss'])
+    assert_equal(state_a['gauss'], state_b['gauss'])


Code review comments:
diff_hunk: @@ -2020,3 +2020,21 @@ def test_broadcast_size_error():
         random.binomial([1, 2], 0.3, size=(2, 1))
     with pytest.raises(ValueError):
         random.binomial([1, 2], [0.3, 0.7], size=(2, 1))
+
+
+def test_randomstate_ctor_old_style_pickle():
+    rs = np.random.RandomState(MT19937(0))
+    rs.standard_normal(1)
+    # Directly call reduce which is used in pickline
path: numpy/random/tests/test_randomstate.py
text: @user1:
```suggestion
    # Directly call reduce which is used in pickling
```

@author:
Fixed, thanks.

diff_hunk: @@ -2695,3 +2695,16 @@ def test_contig_req_out(dist, order, dtype):
     assert variates is out
     variates = dist(out=out, dtype=dtype, size=out.shape)
     assert variates is out
+
+
+def test_generator_ctor_old_style_pickle():
+    rg = np.random.Generator(np.random.PCG64DXSM(0))
+    rg.standard_normal(1)
+    # Directly call reduce which is used in pickline
path: numpy/random/tests/test_generator_mt19937.py
text: @user1:
```suggestion
    # Directly call reduce which is used in pickling
```

@author:
Fixed, thanks.