MAINT: Update typing information

numpy · Jul 15, 2022 · c5c79cd · c5c79cd
1 parent 8147a66
commit c5c79cd
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 2 deletions.
diff --git a/doc/release/upcoming_changes/21976.new_feature.rst b/doc/release/upcoming_changes/21976.new_feature.rst
@@ -6,7 +6,9 @@ The function ``set_bit_generator`` allows the default bit generator to be
 replaced with a user-provided bit generator. This function has been introduced to
 provide a method allowing seemless integration of a high-quality, modern bit generator
 in new code with existing code that makes use of the singleton-provided random
-variate generating functions.
+variate generating functions. The companion function ``get_bit_generator`` returns the current bit generator
+being used by the singleton ``RandomState``. This is provided to simplify restoring
+the original source of randomness if required.
 
 The preferred method to generate reproducible random numbers is to use a modern
 bit generator in an instance of ``Generator``. The function ``default_rng``
@@ -18,8 +20,12 @@ simplifies instantization.
 The same bit generator can then shared with the singleton instance so that
 calling functions in the ``random`` module will use the same bit generator.
 
+    >>> orig_bit_gen = np.random.get_bit_generator()
     >>> np.random.set_bit_generator(rg.bit_generator)
     >>> np.random.normal()
 
 The swap is permanent (until reversed) and so any call to functions
-in the ``random`` module will use the new bit generator.
+in the ``random`` module will use the new bit generator. The orginal
+can be restored if required for code to run correctly.
+
+    >>> np.random.set_bit_generator(orig_bit_gen)
diff --git a/numpy/random/mtrand.pyi b/numpy/random/mtrand.pyi
@@ -562,3 +562,9 @@ zipf = _rand.zipf
 # Two legacy that are trivial wrappers around random_sample
 sample = _rand.random_sample
 ranf = _rand.random_sample
+
+def set_bit_generator(bitgen: BitGenerator) -> None:
+    ...
+
+def get_bit_generator() -> BitGenerator:
+    ...
diff --git a/numpy/typing/tests/data/pass/random.py b/numpy/typing/tests/data/pass/random.py
@@ -1494,3 +1494,6 @@
 random_st.tomaxint()
 random_st.tomaxint(1)
 random_st.tomaxint((1,))
+
+np.random.set_bit_generator(SEED_PCG64)
+np.random.get_bit_generator()
diff --git a/numpy/typing/tests/data/reveal/random.pyi b/numpy/typing/tests/data/reveal/random.pyi
@@ -1537,3 +1537,6 @@ reveal_type(random_st.random_sample(size=(1, 2)))  # E: ndarray[Any, dtype[float
 reveal_type(random_st.tomaxint())  # E: int
 reveal_type(random_st.tomaxint(1))  # E: ndarray[Any, dtype[{int_}]]
 reveal_type(random_st.tomaxint((1,)))  # E: ndarray[Any, dtype[{int_}]]
+
+reveal_type(np.random.set_bit_generator(pcg64))  # E: BitGenerator
+reveal_type(np.random.get_bit_generator())  # E: None