From b2fd5e1480cdbb71e207aa1ab08f0275929d3dc4 Mon Sep 17 00:00:00 2001
From: ViciousEagle03 <piyushsharma04321@gmail.com>
Date: Wed, 21 Aug 2024 20:41:55 +0530
Subject: [PATCH] Add mask as a validator property

---
 ndcube/asdf/converters/ndcube_converter.py            | 9 +++++++--
 ndcube/asdf/converters/tests/test_ndcube_converter.py | 2 +-
 ndcube/asdf/resources/schemas/ndcube-0.1.0.yaml       | 2 ++
 ndcube/conftest.py                                    | 9 +++++++--
 4 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/ndcube/asdf/converters/ndcube_converter.py b/ndcube/asdf/converters/ndcube_converter.py
index 34d372019..c88b8bdff 100644
--- a/ndcube/asdf/converters/ndcube_converter.py
+++ b/ndcube/asdf/converters/ndcube_converter.py
@@ -10,7 +10,10 @@ class NDCubeConverter(Converter):
     def from_yaml_tree(self, node, tag, ctx):
         from ndcube.ndcube import NDCube
 
-        ndcube = NDCube(node["data"], node["wcs"], meta=node.get("meta"))
+        ndcube = NDCube(node["data"],
+                        node["wcs"],
+                        meta = node.get("meta"),
+                        mask = node.get("mask"))
         if "extra_coords" in node:
             ndcube._extra_coords = node["extra_coords"]
         if "global_coords" in node:
@@ -41,8 +44,10 @@ def to_yaml_tree(self, ndcube, tag, ctx):
         node["extra_coords"] = ndcube.extra_coords
         node["global_coords"] = ndcube.global_coords
         node["meta"] = ndcube.meta
+        if ndcube.mask is not None:
+            node["mask"] = ndcube.mask
 
-        attributes = ['uncertainty', 'mask', 'unit']
+        attributes = ['uncertainty', 'unit']
         for attr in attributes:
             if getattr(ndcube, attr) is not None:
                 warnings.warn(f"Attribute '{attr}' is present but not being saved in ASDF serialization.", UserWarning)
diff --git a/ndcube/asdf/converters/tests/test_ndcube_converter.py b/ndcube/asdf/converters/tests/test_ndcube_converter.py
index 19feefa1b..ad5df5ee7 100644
--- a/ndcube/asdf/converters/tests/test_ndcube_converter.py
+++ b/ndcube/asdf/converters/tests/test_ndcube_converter.py
@@ -8,7 +8,7 @@
 from ndcube.tests.helpers import assert_cubes_equal
 
 
-@pytest.mark.parametrize("ndc",[("ndcube_gwcs_2d_ln_lt"),
+@pytest.mark.parametrize("ndc",[("ndcube_gwcs_2d_ln_lt_mask"),
                                 ("ndcube_gwcs_3d_ln_lt_l"),
                                 ("ndcube_gwcs_3d_ln_lt_l_ec_dropped_dim"),
                                 ("ndcube_gwcs_3d_ln_lt_l_ec_q_t_gc"),
diff --git a/ndcube/asdf/resources/schemas/ndcube-0.1.0.yaml b/ndcube/asdf/resources/schemas/ndcube-0.1.0.yaml
index db10a488e..5529ccbfa 100644
--- a/ndcube/asdf/resources/schemas/ndcube-0.1.0.yaml
+++ b/ndcube/asdf/resources/schemas/ndcube-0.1.0.yaml
@@ -21,6 +21,8 @@ properties:
     tag: "tag:sunpy.org:ndcube/global_coords/globalcoords-0.*"
   meta:
     type: object
+  mask:
+    type: object
 
 required: [data, wcs]
 additionalProperties: true
diff --git a/ndcube/conftest.py b/ndcube/conftest.py
index 56b8a0165..6d0dcf9c5 100644
--- a/ndcube/conftest.py
+++ b/ndcube/conftest.py
@@ -502,10 +502,15 @@ def ndcube_gwcs_3d_ln_lt_l_ec_q_t_gc(gwcs_3d_lt_ln_l):
     return cube
 
 @pytest.fixture
-def ndcube_gwcs_2d_ln_lt(gwcs_2d_lt_ln):
+def ndcube_gwcs_2d_ln_lt_mask(gwcs_2d_lt_ln):
     shape = (10, 12)
     data_cube = data_nd(shape)
-    return NDCube(data_cube, wcs=gwcs_2d_lt_ln)
+    mask = np.zeros(shape, dtype=bool)
+    mask[1, 1] = True
+    mask[2, 0] = True
+    mask[3, 3] = True
+    mask[4:6, :4] = True
+    return NDCube(data_cube, wcs=gwcs_2d_lt_ln, mask=mask)
 
 @pytest.fixture
 def ndcube_4d_ln_l_t_lt(wcs_4d_lt_t_l_ln):