task: review fixes

Author: jharlow-intel

.pylintrc

@@ -99,7 +99,6 @@
     is_patched,
     mkl_random,
     patch_numpy_random,
-    patched_names,
     restore_numpy_random,
 )
 
@@ -155,6 +154,10 @@
     "shuffle",
     "permutation",
     "interfaces",
+    "mkl_random",
+    "patch_numpy_random",
+    "restore_numpy_random",
+    "is_patched",
 ]
 
 del _init_helper

mkl_random/__init__.py

@@ -25,47 +25,29 @@
 
 """Define functions for patching NumPy with MKL-based NumPy interface."""
 
-import warnings
 from contextlib import ContextDecorator
 from threading import Lock, local
 
-import numpy as _np
+import numpy as np
 
 import mkl_random.interfaces.numpy_random as _nrand
 
-_DEFAULT_NAMES = tuple(_nrand.__all__)
-
 
 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._active_names = ()
-        self._patched = ()
+        # make _patched_functions a tuple (immutable)
+        self._patched_functions = tuple(_nrand.__all__)
         self._tls = local()
 
-    def _normalize_names(self, names):
-        if names is None:
-            names = _DEFAULT_NAMES
-        return tuple(names)
-
-    def _validate_module(self, numpy_module):
-        if not hasattr(numpy_module, "random"):
-            raise TypeError(
-                "Expected a numpy-like module with a `.random` attribute."
-            )
-
     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)
+            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:
@@ -79,51 +61,12 @@ def _restore_func(self, name, verbose=False):
         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
-        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(_nrand, name):
-                missing.append(name)
-                continue
-            patchable.append(name)
-
-        if strict and missing:
-            raise AttributeError(
-                "Could not patch these names (missing on numpy.random or "
-                "mkl_random.interfaces.numpy_random): "
-                + ", ".join(str(x) for x in missing)
-            )
-
-        self._numpy_module = numpy_module
-        self._requested_names = names
-        self._active_names = tuple(patchable)
-        self._patched = tuple(patchable)
-
-    def do_patch(
-        self,
-        numpy_module=None,
-        names=None,
-        strict=False,
-        verbose=False,
-    ):
-        if numpy_module is None:
-            numpy_module = _np
-        names = self._normalize_names(names)
-        strict = bool(strict)
+            setattr(np.random, name, val)
 
+    def do_patch(self, verbose=False):
         with self._lock:
             local_count = getattr(self._tls, "local_count", 0)
             if self._patch_count == 0:
-                self._initialize_patch(numpy_module, names, strict)
                 if verbose:
                     print(
                         "Now patching NumPy random submodule with mkl_random "
@@ -133,19 +76,8 @@ def do_patch(
                         "Please direct bug reports to "
                         "https://github.com/IntelPython/mkl_random"
                     )
-                for name in self._active_names:
-                    self._register_func(name, getattr(_nrand, name))
-            else:
-                if self._numpy_module is not numpy_module:
-                    raise RuntimeError(
-                        "Already patched a different numpy module; "
-                        "call restore() first."
-                    )
-                if names != self._requested_names:
-                    raise RuntimeError(
-                        "Already patched with a different names set; "
-                        "call restore() first."
-                    )
+                for f in self._patched_functions:
+                    self._register_func(f, getattr(_nrand, f))
             self._patch_count += 1
             self._tls.local_count = local_count + 1
 
@@ -154,77 +86,47 @@ def do_restore(self, verbose=False):
             local_count = getattr(self._tls, "local_count", 0)
             if local_count <= 0:
                 if verbose:
-                    warnings.warn(
+                    print(
                         "Warning: restore_numpy_random called more times than "
-                        "patch_numpy_random in this thread.",
-                        stacklevel=2,
+                        "patch_numpy_random in this thread."
                     )
                 return
-
-            self._tls.local_count = local_count - 1
+            self._tls.local_count -= 1
             self._patch_count -= 1
             if self._patch_count == 0:
                 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._active_names = ()
-                self._patched = ()
 
     def is_patched(self):
         with self._lock:
             return self._patch_count > 0
 
-    def patched_names(self):
-        with self._lock:
-            return list(self._patched)
-
 
 _patch = _GlobalPatch()
 
 
-def patch_numpy_random(
-    numpy_module=None,
-    names=None,
-    strict=False,
-    verbose=False,
-):
+def patch_numpy_random(verbose=False):
     """
-    Patch NumPy's random submodule with mkl_random's NumPy interface.
+    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.
+        print message when starting the patching process.
+
+    Notes
+    -----
+    This function uses reference-counted semantics. Each call increments a
+    global patch counter. Restoration requires a matching number of calls
+    between `patch_numpy_random` and `restore_numpy_random`.
+
+    In multi-threaded programs, prefer the `mkl_random` context manager.
 
-    Examples
-    --------
-    >>> import numpy as np
-    >>> import mkl_random
-    >>> mkl_random.is_patched()
-    False
-    >>> mkl_random.patch_numpy_random(np)
-    >>> mkl_random.is_patched()
-    True
-    >>> mkl_random.restore()
-    >>> mkl_random.is_patched()
-    False
     """
-    _patch.do_patch(
-        numpy_module=numpy_module,
-        names=names,
-        strict=bool(strict),
-        verbose=bool(verbose),
-    )
+    _patch.do_patch(verbose=verbose)
 
 
 def restore_numpy_random(verbose=False):
@@ -234,45 +136,47 @@ def restore_numpy_random(verbose=False):
     Parameters
     ----------
     verbose : bool, optional
-        Print message when starting restoration process.
+        print message when starting restoration process.
+
+    Notes
+    -----
+    This function uses reference-counted semantics. Each call decrements a
+    global patch counter. Restoration requires a matching number of calls
+    between `patch_numpy_random` and `restore_numpy_random`.
+
+    In multi-threaded programs, prefer the `mkl_random` context manager.
+
     """
-    _patch.do_restore(verbose=bool(verbose))
+    _patch.do_restore(verbose=verbose)
 
 
 def is_patched():
-    """Return whether NumPy has been patched with mkl_random."""
+    """Return True if NumPy's random sm is currently patched by mkl_random."""
     return _patch.is_patched()
 
 
-def patched_names():
-    """Return names actually patched in `numpy.random`."""
-    return _patch.patched_names()
-
-
 class mkl_random(ContextDecorator):
     """
     Context manager and decorator to temporarily patch NumPy random submodule
     with MKL-based implementations.
 
     Examples
     --------
-    >>> import numpy as np
     >>> import mkl_random
-    >>> with mkl_random.mkl_random(np):
-    ...     x = np.random.normal(size=10)
-    """
+    >>> mkl_random.is_patched()
+    # False
+
+    >>> with mkl_random.mkl_random():  # Enable mkl_random in NumPy
+    >>>     print(mkl_random.is_patched())
+    # True
 
-    def __init__(self, numpy_module=None, names=None, strict=False):
-        self._numpy_module = numpy_module
-        self._names = names
-        self._strict = strict
+    >>> mkl_random.is_patched()
+    # False
+
+    """
 
     def __enter__(self):
-        patch_numpy_random(
-            numpy_module=self._numpy_module,
-            names=self._names,
-            strict=self._strict,
-        )
+        patch_numpy_random()
         return self
 
     def __exit__(self, *exc):