Checking if an indexer is a scalar or not is now more robust

Author: pawel-wolff

src/earthkit/data/indexing/tensor.py

@@ -422,14 +422,13 @@ def _prepare_tensor_data(self, source_to_array_func, index=None):
                 # * `field_shape` does lose the dimension `dim` if `dim` is a field dimension
                 current_field_shape = []
                 for n, _idx in zip(self.field_shape, index):
-                    if isinstance(_idx, int):
-                        # simply, ignore this index
+                    _sizes = np.arange(n)[_idx].shape
+                    if len(_sizes) == 0:
+                        # _idx is a scalar indexer, and thus we ignore it
                         continue
-                    if isinstance(_idx, slice):
-                        _size = len(range(n)[_idx])
-                    else:
-                        # _idx must be an iterable of integers
-                        _size = len(np.arange(n)[_idx])
+                    # _idx is a slice, an array of int's or a boolean mask
+                    # get the size of the selection made by the indexer _idx
+                    (_size,) = _sizes
                     current_field_shape.append(_size)
                 current_field_shape = tuple(current_field_shape)
 
@@ -485,7 +484,7 @@ def _subset(self, indexes):
         user_indexes = []
 
         for s, c in zip(indexes, self._user_shape):
-            lst = np.array(list(range(c)))[s].tolist()
+            lst = np.arange(c)[s].tolist()
             if not isinstance(lst, list):
                 lst = [lst]
             user_coords.append(lst)
@@ -687,7 +686,7 @@ def _subset(self, indexes):
         user_indexes = []
 
         for s, c in zip(indexes, self._user_shape):
-            lst = np.array(list(range(c)))[s].tolist()
+            lst = np.arange(c)[s].tolist()
             if not isinstance(lst, list):
                 lst = [lst]
             user_icoords.append(lst)

src/earthkit/data/utils/array.py

@@ -120,7 +120,7 @@ def outer_indexing(v, indices):
         v = v[_1d_index]
         v_ndim = v.ndim
         if v_ndim == ndim:
-            full_slices = full_slices + (slice(None, None),)
+            full_slices = full_slices + (slice(None),)
         else:
             # the current dimension has collapsed
             ndim = v_ndim