improve mask access

fabm-model · Oct 21, 2024 · 3e6ff12 · 3e6ff12
1 parent 9ed3b9e
commit 3e6ff12
Showing 1 changed file with 9 additions and 6 deletions.
diff --git a/src/pyfabm/__init__.py b/src/pyfabm/__init__.py
@@ -1087,10 +1087,15 @@ def link_mask(self, *masks: np.ndarray):
         self._mask = masks
         self.fabm.set_mask(self.pmodel, *self._mask)
 
-    def _get_mask(self) -> Union[np.ndarray, Sequence[np.ndarray]]:
-        return self._mask[0] if len(self._mask) == 1 else self._mask
-
-    def _set_mask(self, values: Union[npt.ArrayLike, Sequence[npt.ArrayLike]]):
+    @property
+    def mask(self) -> Union[np.ndarray, Sequence[np.ndarray], None]:
+        mask = self._mask
+        if mask is not None and len(mask) == 1:
+            mask = mask[0]
+        return mask
+
+    @mask.setter
+    def mask(self, values: Union[npt.ArrayLike, Sequence[npt.ArrayLike]]):
         if self.fabm.mask_type == 1:
             values = (values,)
         if len(values) != self.fabm.mask_type:
@@ -1104,8 +1109,6 @@ def _set_mask(self, values: Union[npt.ArrayLike, Sequence[npt.ArrayLike]]):
             if value is not mask:
                 mask[...] = value
 
-    mask = property(_get_mask, _set_mask)
-
     def link_bottom_index(self, indices: np.ndarray):
         if not self.fabm.variable_bottom_index:
             raise FABMException(