Merge pull request #21 from jcmgray/mimicfix

fix NumpyMimic special attribute access (#20)

autoray/autoray.py

@@ -1448,7 +1448,25 @@ def binary_dispatcher(args, kwargs):
 
 # --------------- object to act as drop-in replace for numpy ---------------- #
 
-_partial_functions = {}
+
+def _get_mimic_function_or_attribute(self, fn):
+    # respect all 'dunder' special methods and attributes
+    if (fn[:2] == "__") and (fn[-2:] == "__"):
+        return object.__getattribute__(self, fn)
+
+    # look out for certain submodules which are not functions
+    if fn == "linalg":
+        return NumpyMimic("linalg")
+
+    if fn == "random":
+        return NumpyMimic("random")
+
+    # if this is the e.g. linalg mimic, preprend 'linalg.'
+    submod = object.__getattribute__(self, "submodule")
+    if submod is not None:
+        fn = ".".join((submod, fn))
+
+    return functools.partial(do, fn)
 
 
 class NumpyMimic:
@@ -1457,23 +1475,13 @@ class NumpyMimic:
     def __init__(self, submodule=None):
         self.submodule = submodule
 
-    def __getattribute__(self, fn):
-        # look out for certain submodules which are not functions
-        if fn == "linalg":
-            return numpy_linalg
-        if fn == "random":
-            return numpy_random
-
-        # if this is the e.g. linalg mimic, preprend 'linalg.'
-        submod = object.__getattribute__(self, "submodule")
-        if submod is not None:
-            fn = ".".join((submod, fn))
-
-        # cache the correct partial function
+    def __getattribute__(self, attr):
+        # cache the correct partial function (or special method/attribute)
+        d = object.__getattribute__(self, "__dict__")
         try:
-            pfn = _partial_functions[fn]
+            pfn = d[attr]
         except KeyError:
-            pfn = _partial_functions[fn] = functools.partial(do, fn)
+            pfn = d[attr] = _get_mimic_function_or_attribute(self, attr)
 
         return pfn
 
@@ -1483,8 +1491,6 @@ def __repr__():
 
 
 numpy = NumpyMimic()
-numpy_linalg = NumpyMimic("linalg")
-numpy_random = NumpyMimic("random")
 
 
 # --------------------------------------------------------------------------- #

tests/test_autoray.py

@@ -158,6 +158,19 @@ def modified_gram_schmidt_np_mimic(X):
     return np.stack(Q, axis=0)
 
 
+def test_numpy_mimic_dunder_methods():
+    from abc import ABC
+    from autoray import numpy as np
+
+    class Base(ABC):
+        pass
+
+    assert isinstance(np, object)
+    assert not isinstance(np, Base)
+    print(np)
+    dir(np)
+
+
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_mgs_np_mimic(backend):
     if backend == "sparse":