python - Can I create a local numpy random seed?

Question

There is a function, foo, that uses the np.random functionality. I want to control the seed that foo uses, but without actually changing the function itself. How do I do this?

Essentially I want something like this:

bar() # should have normal seed
with np.random.seed(0): # Doesn't work
    foo()
bar() # should have normal seed

---

Solutions like this:

rng = random.Random(42)
number = rng.randint(10, 20)

doesn't work in this case, as I don't have access to the inner workings of foo (or am I missing something??).

Answer 1

You could keep the global random state in a temporary variable and reset it once your function is done:

import contextlib
import numpy as np

@contextlib.contextmanager
def temp_seed(seed):
    state = np.random.get_state()
    np.random.seed(seed)
    try:
        yield
    finally:
        np.random.set_state(state)

Demo:

>>> np.random.seed(0)
>>> np.random.randn(3)
array([1.76405235, 0.40015721, 0.97873798])
>>> np.random.randn(3)
array([ 2.2408932 ,  1.86755799, -0.97727788])

>>> np.random.seed(0)
>>> np.random.randn(3)
array([1.76405235, 0.40015721, 0.97873798])
>>> with temp_seed(5):
...     np.random.randn(3)                                                                                        
array([ 0.44122749, -0.33087015,  2.43077119])
>>> np.random.randn(3)
array([ 2.2408932 ,  1.86755799, -0.97727788])