fix: preserve masked array in view

Author: flying-sheep

src/anndata/_core/views.py

@@ -195,6 +195,25 @@ def toarray(self) -> np.ndarray:
         return self.copy()
 
 
+class MaskedArrayView(_SetItemMixin, np.ma.MaskedArray):
+    def __new__(
+        cls,
+        input_array: Sequence[Any],
+        view_args: ViewArgs | None = None,
+    ):
+        arr = np.ma.asarray(input_array).view(cls)
+
+        if view_args is not None:
+            view_args = ElementRef(*view_args)
+        arr._view_args = view_args
+        return arr
+
+    def __array_finalize__(self, obj: np.ndarray | None):
+        super().__array_finalize__(obj)
+        if obj is not None:
+            self._view_args = getattr(obj, "_view_args", None)
+
+
 # Extends DaskArray
 # Calls parent __new__ constructor since
 # even calling astype on a dask array
@@ -341,6 +360,11 @@ def as_view_array(array, view_args):
     return ArrayView(array, view_args=view_args)
 
 
+@as_view.register(np.ma.MaskedArray)
+def as_view_masked_array(array, view_args):
+    return MaskedArrayView(array, view_args=view_args)
+
+
 @as_view.register(DaskArray)
 def as_view_dask_array(array, view_args):
     return DaskArrayView(array, view_args=view_args)

tests/test_views.py

@@ -29,6 +29,7 @@
 )
 from anndata._core.views import (
     ArrayView,
+    MaskedArrayView,
     SparseCSCArrayView,
     SparseCSCMatrixView,
     SparseCSRArrayView,
@@ -709,6 +710,29 @@ def test_view_retains_ndarray_subclass():
     assert view.obsm["foo"].shape == (5, 5)
 
 
+def test_view_of_masked_array():
+    mask = np.zeros((10, 5), dtype=bool)
+    mask[0, 0] = True
+    data = np.ma.MaskedArray(np.arange(50.0).reshape(10, 5), mask=mask)
+
+    adata = ad.AnnData(np.zeros((10, 10)), obsm={"masked": data})
+    view = adata[:5, :]
+
+    masked_view = view.obsm["masked"]
+    assert isinstance(masked_view, MaskedArrayView)
+    assert isinstance(masked_view, np.ma.MaskedArray)
+    assert masked_view.mask[0, 0]
+    assert not masked_view.mask[1, 0]
+    assert np.ma.is_masked(masked_view)
+
+    with pytest.warns(ImplicitModificationWarning, match=r"obsm"):
+        masked_view[1, 0] = 1234.0
+    assert view.obsm["masked"][1, 0] == 1234.0
+    # Original is untouched: writing to a view actualizes it as a copy.
+    assert adata.obsm["masked"][1, 0] == 5.0
+    assert adata.obsm["masked"].mask[0, 0]
+
+
 def test_modify_uns_in_copy():
     # https://github.com/scverse/anndata/issues/571
     adata = ad.AnnData(np.ones((5, 5)), uns={"parent": {"key": "value"}})