From 5b7b2de1e214fc975718da327c41b3e6307f0b1e Mon Sep 17 00:00:00 2001
From: Baudouin Raoult <baudouin.raoult@ecmwf.int>
Date: Wed, 13 Mar 2024 10:21:08 +0000
Subject: [PATCH] mask indexing

---
 ecml_tools/data.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/ecml_tools/data.py b/ecml_tools/data.py
index 5910b13..316bf0c 100644
--- a/ecml_tools/data.py
+++ b/ecml_tools/data.py
@@ -885,14 +885,26 @@ def longitudes(self):
 
     @debug_indexing
     def __getitem__(self, index):
+        if isinstance(index, tuple):
+            return self._get_tuple(index)
 
         result = self.forward[index]
-        print(index, result.shape)
         # We don't support subsetting the grid values
         assert result.shape[-1] == len(self.mask), (result.shape, len(self.mask))
 
         return result[..., self.mask]
 
+    @debug_indexing
+    @expand_list_indexing
+    def _get_tuple(self, index):
+        index, changes = index_to_slices(index, self.shape)
+        index, previous = update_tuple(index, self.axis, slice(None))
+        result = self.forward[index]
+        result = result[..., self.mask]
+        result = result[..., previous]
+        result = apply_index_to_slices_changes(result, changes)
+        return result
+
 
 class Thinning(Masked):