From 085e559d42eb3066d1d5756c9eb05aea62839a03 Mon Sep 17 00:00:00 2001
From: ayushnag <35325113+ayushnag@users.noreply.github.com>
Date: Mon, 15 Jul 2024 10:15:07 -0700
Subject: [PATCH] ufunc isnan fix

---
 virtualizarr/manifests/array.py     |  4 +++-
 virtualizarr/manifests/array_api.py | 10 +++++-----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/virtualizarr/manifests/array.py b/virtualizarr/manifests/array.py
index a0983de..e15cf7d 100644
--- a/virtualizarr/manifests/array.py
+++ b/virtualizarr/manifests/array.py
@@ -5,7 +5,7 @@
 
 from ..kerchunk import KerchunkArrRefs
 from ..zarr import ZArray
-from .array_api import MANIFESTARRAY_HANDLED_ARRAY_FUNCTIONS
+from .array_api import MANIFESTARRAY_HANDLED_ARRAY_FUNCTIONS, _isnan
 from .manifest import ChunkManifest
 
 
@@ -127,6 +127,8 @@ def __array_function__(self, func, types, args, kwargs) -> Any:
 
     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs) -> Any:
         """We have to define this in order to convince xarray that this class is a duckarray, even though we will never support ufuncs."""
+        if ufunc == np.isnan:
+            return _isnan(self.shape)
         return NotImplemented
 
     def __array__(self) -> np.ndarray:
diff --git a/virtualizarr/manifests/array_api.py b/virtualizarr/manifests/array_api.py
index 0ecdc02..0960697 100644
--- a/virtualizarr/manifests/array_api.py
+++ b/virtualizarr/manifests/array_api.py
@@ -356,8 +356,8 @@ def isnan(x: "ManifestArray", /) -> np.ndarray:
 
     Only implemented to get past some checks deep inside xarray, see https://github.com/TomNicholas/VirtualiZarr/issues/29.
     """
-    return np.full(
-        shape=x.shape,
-        fill_value=False,
-        dtype=np.dtype(bool),
-    )
+    return _isnan(x.shape)
+
+
+def _isnan(shape: tuple):
+    return np.full(shape=shape, fill_value=False, dtype=np.dtype(bool))