fix: review

Author: jharlow-intel

mkl_random/__init__.py

@@ -99,8 +99,10 @@
     is_patched,
     mkl_random,
     monkey_patch,
+    patch_numpy_random,
     patched_names,
     restore,
+    restore_numpy_random,
     use_in_numpy,
 )
 

mkl_random/src/_patch.pyx mkl_random/_patch.pymkl_random/src/_patch.pyx renamed to mkl_random/_patch.py

@@ -23,17 +23,7 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-# distutils: language = c
-# cython: language_level=3
-
-"""
-Patch NumPy's `numpy.random` symbols to use mkl_random implementations.
-
-This is attribute-level monkey patching. It can replace legacy APIs like
-`numpy.random.RandomState` and global distribution functions, but it does not
-replace NumPy's `Generator`/`default_rng()` unless mkl_random provides fully
-compatible replacements.
-"""
+"""Define functions for patching NumPy with MKL-based NumPy interface."""
 
 from contextlib import ContextDecorator
 from threading import Lock, local
@@ -43,24 +33,21 @@
 from . import mklrand as _mr
 
 
-cdef tuple _DEFAULT_NAMES = (
+_DEFAULT_NAMES = (
     # Legacy seeding / state
     "seed",
     "get_state",
     "set_state",
     "RandomState",
-
     # Common global sampling helpers
     "random",
     "random_sample",
     "sample",
     "rand",
     "randn",
     "bytes",
-
     # Integers
     "randint",
-
     # Common distributions (only patched if present on both sides)
     "standard_normal",
     "normal",
@@ -82,7 +69,6 @@
     "wald",
     "weibull",
     "zipf",
-
     # Permutations / choices
     "choice",
     "permutation",
@@ -94,9 +80,11 @@ class _GlobalPatch:
     def __init__(self):
         self._lock = Lock()
         self._patch_count = 0
+        self._restore_dict = {}
+        self._patched_functions = tuple(_DEFAULT_NAMES)
         self._numpy_module = None
         self._requested_names = None
-        self._originals = {}
+        self._active_names = ()
         self._patched = ()
         self._tls = local()
 
@@ -111,22 +99,44 @@ def _validate_module(self, numpy_module):
                 "Expected a numpy-like module with a `.random` attribute."
             )
 
-    def _apply_patch(self, numpy_module, names, strict):
+    def _register_func(self, name, func):
+        if name not in self._patched_functions:
+            raise ValueError(f"{name} not an mkl_random function.")
+        np_random = self._numpy_module.random
+        if name not in self._restore_dict:
+            self._restore_dict[name] = getattr(np_random, name)
+        setattr(np_random, name, func)
+
+    def _restore_func(self, name, verbose=False):
+        if name not in self._patched_functions:
+            raise ValueError(f"{name} not an mkl_random function.")
+        try:
+            val = self._restore_dict[name]
+        except KeyError:
+            if verbose:
+                print(f"failed to restore {name}")
+            return
+        else:
+            if verbose:
+                print(f"found and restoring {name}...")
+            np_random = self._numpy_module.random
+            setattr(np_random, name, val)
+
+    def _initialize_patch(self, numpy_module, names, strict):
+        self._validate_module(numpy_module)
         np_random = numpy_module.random
-        originals = {}
-        patched = []
         missing = []
+        patchable = []
         for name in names:
+            if name not in self._patched_functions:
+                missing.append(name)
+                continue
             if not hasattr(np_random, name) or not hasattr(_mr, name):
                 missing.append(name)
                 continue
-            originals[name] = getattr(np_random, name)
-            setattr(np_random, name, getattr(_mr, name))
-            patched.append(name)
+            patchable.append(name)
 
         if strict and missing:
-            for name, value in originals.items():
-                setattr(np_random, name, value)
             raise AttributeError(
                 "Could not patch these names (missing on numpy.random or "
                 "mkl_random.mklrand): "
@@ -135,8 +145,8 @@ def _apply_patch(self, numpy_module, names, strict):
 
         self._numpy_module = numpy_module
         self._requested_names = names
-        self._originals = originals
-        self._patched = tuple(patched)
+        self._active_names = tuple(patchable)
+        self._patched = tuple(patchable)
 
     def do_patch(
         self,
@@ -148,13 +158,23 @@ def do_patch(
         if numpy_module is None:
             numpy_module = _np
         names = self._normalize_names(names)
-        self._validate_module(numpy_module)
         strict = bool(strict)
 
         with self._lock:
             local_count = getattr(self._tls, "local_count", 0)
             if self._patch_count == 0:
-                self._apply_patch(numpy_module, names, strict)
+                self._initialize_patch(numpy_module, names, strict)
+                if verbose:
+                    print(
+                        "Now patching NumPy random submodule with mkl_random "
+                        "NumPy interface."
+                    )
+                    print(
+                        "Please direct bug reports to "
+                        "https://github.com/IntelPython/mkl_random"
+                    )
+                for name in self._active_names:
+                    self._register_func(name, getattr(_mr, name))
             else:
                 if self._numpy_module is not numpy_module:
                     raise RuntimeError(
@@ -175,20 +195,22 @@ def do_restore(self, verbose=False):
             if local_count <= 0:
                 if verbose:
                     print(
-                        "Warning: restore called more times than monkey_patch "
-                        "in this thread."
+                        "Warning: restore_numpy_random called more times than "
+                        "patch_numpy_random in this thread."
                     )
                 return
 
             self._tls.local_count = local_count - 1
             self._patch_count -= 1
             if self._patch_count == 0:
-                np_random = self._numpy_module.random
-                for name, value in self._originals.items():
-                    setattr(np_random, name, value)
+                if verbose:
+                    print("Now restoring original NumPy random submodule.")
+                for name in tuple(self._restore_dict):
+                    self._restore_func(name, verbose=verbose)
+                self._restore_dict.clear()
                 self._numpy_module = None
                 self._requested_names = None
-                self._originals = {}
+                self._active_names = ()
                 self._patched = ()
 
     def is_patched(self):
@@ -203,18 +225,33 @@ def patched_names(self):
 _patch = _GlobalPatch()
 
 
-def monkey_patch(numpy_module=None, names=None, strict=False, verbose=False):
+def patch_numpy_random(
+    numpy_module=None,
+    names=None,
+    strict=False,
+    verbose=False,
+):
     """
-    Enables using mkl_random in the given NumPy module by patching
-    `numpy.random`.
+    Patch NumPy's random submodule with mkl_random's NumPy interface.
+
+    Parameters
+    ----------
+    numpy_module : module, optional
+        NumPy-like module to patch. Defaults to imported NumPy.
+    names : iterable[str], optional
+        Attributes under `numpy_module.random` to patch.
+    strict : bool, optional
+        Raise if any requested symbol cannot be patched.
+    verbose : bool, optional
+        Print messages when starting the patching process.
 
     Examples
     --------
     >>> import numpy as np
     >>> import mkl_random
     >>> mkl_random.is_patched()
     False
-    >>> mkl_random.monkey_patch(np)
+    >>> mkl_random.patch_numpy_random(np)
     >>> mkl_random.is_patched()
     True
     >>> mkl_random.restore()
@@ -229,11 +266,31 @@ def monkey_patch(numpy_module=None, names=None, strict=False, verbose=False):
     )
 
 
-def use_in_numpy(numpy_module=None, names=None, strict=False, verbose=False):
+def restore_numpy_random(verbose=False):
     """
-    Backward-compatible alias for monkey_patch().
+    Restore NumPy's random submodule to its original implementations.
+
+    Parameters
+    ----------
+    verbose : bool, optional
+        Print message when starting restoration process.
     """
-    monkey_patch(
+    _patch.do_restore(verbose=bool(verbose))
+
+
+def monkey_patch(numpy_module=None, names=None, strict=False, verbose=False):
+    """Backward-compatible alias for patch_numpy_random()."""
+    patch_numpy_random(
+        numpy_module=numpy_module,
+        names=names,
+        strict=strict,
+        verbose=verbose,
+    )
+
+
+def use_in_numpy(numpy_module=None, names=None, strict=False, verbose=False):
+    """Backward-compatible alias for patch_numpy_random()."""
+    patch_numpy_random(
         numpy_module=numpy_module,
         names=names,
         strict=strict,
@@ -242,11 +299,8 @@ def use_in_numpy(numpy_module=None, names=None, strict=False, verbose=False):
 
 
 def restore(verbose=False):
-    """
-    Disables using mkl_random in NumPy by restoring the original
-    `numpy.random` symbols.
-    """
-    _patch.do_restore(verbose=bool(verbose))
+    """Backward-compatible alias for restore_numpy_random()."""
+    restore_numpy_random(verbose=verbose)
 
 
 def is_patched():
@@ -265,7 +319,8 @@ def patched_names():
 
 class mkl_random(ContextDecorator):
     """
-    Context manager and decorator to temporarily patch NumPy's `numpy.random`.
+    Context manager and decorator to temporarily patch NumPy random submodule
+    with MKL-based implementations.
 
     Examples
     --------
@@ -274,19 +329,20 @@ class mkl_random(ContextDecorator):
     >>> with mkl_random.mkl_random(np):
     ...     x = np.random.normal(size=10)
     """
+
     def __init__(self, numpy_module=None, names=None, strict=False):
         self._numpy_module = numpy_module
         self._names = names
         self._strict = strict
 
     def __enter__(self):
-        monkey_patch(
+        patch_numpy_random(
             numpy_module=self._numpy_module,
             names=self._names,
             strict=self._strict,
         )
         return self
 
     def __exit__(self, *exc):
-        restore()
+        restore_numpy_random()
         return False

mkl_random/tests/test_patch.py

@@ -133,19 +133,20 @@ def test_patch_redundant_patching():
     orig_normal = np.random.normal
     assert not mkl_random.is_patched()
 
-    mkl_random.monkey_patch(np)
-    mkl_random.monkey_patch(np)
-
-    assert mkl_random.is_patched()
-    assert np.random.normal is mkl_random.normal
-
-    mkl_random.restore()
-    assert mkl_random.is_patched()
-    assert np.random.normal is mkl_random.normal
-
-    mkl_random.restore()
-    assert not mkl_random.is_patched()
-    assert np.random.normal is orig_normal
+    try:
+        mkl_random.monkey_patch(np)
+        mkl_random.monkey_patch(np)
+        assert mkl_random.is_patched()
+        assert np.random.normal is mkl_random.mklrand.normal
+        mkl_random.restore()
+        assert mkl_random.is_patched()
+        assert np.random.normal is mkl_random.mklrand.normal
+        mkl_random.restore()
+        assert not mkl_random.is_patched()
+        assert np.random.normal is orig_normal
+    finally:
+        while mkl_random.is_patched():
+            mkl_random.restore()
 
 
 def test_patch_reentrant():

setup.py

@@ -93,13 +93,6 @@ def extensions():
             define_macros=defs + [("NDEBUG", None)],
             language="c++",
         ),
-        Extension(
-            "mkl_random._patch",
-            sources=[join("mkl_random", "src", "_patch.pyx")],
-            include_dirs=[np.get_include()],
-            define_macros=defs + [("NDEBUG", None)],
-            language="c",
-        ),
     ]
 
     return exts