Fix indexing for lazy expressions, and allow use of None in getitem

Author: Luke Shaw

src/blosc2/lazyexpr.py

@@ -2440,16 +2440,16 @@ def _compute_expr(self, item, kwargs):  # noqa: C901
                     if len(self._where_args) == 1:
                         # We have a single argument
                         where_x = self._where_args["_where_x"]
-                        return where_x[:][lazy_expr]
+                        return (where_x[:][lazy_expr])[item]
                     if len(self._where_args) == 2:
                         # We have two arguments
                         where_x = self._where_args["_where_x"]
                         where_y = self._where_args["_where_y"]
-                        return np.where(lazy_expr, where_x, where_y)
+                        return np.where(lazy_expr, where_x, where_y)[item]
                 if hasattr(self, "_output"):
                     # This is not exactly optimized, but it works for now
-                    self._output[:] = lazy_expr
-                return lazy_expr
+                    self._output[:] = lazy_expr[item]
+                return lazy_expr[item]
 
             return chunked_eval(lazy_expr.expression, lazy_expr.operands, item, **kwargs)
 

src/blosc2/ndarray.py

@@ -61,14 +61,15 @@ def make_key_hashable(key):
     else:
         return key
 
-
 def process_key(key, shape):
-    if key is None:
-        key = tuple(slice(None) for _ in range(len(shape)))
     key = ndindex.ndindex(key).expand(shape).raw
-    mask = tuple(isinstance(k, int) for k in key)
-    key = tuple(k if isinstance(k, slice) else slice(k, k + 1, None) for k in key)
-    return key, mask
+    key_ = () 
+    mask = () #get integer indices where dimension collapses
+    for k in key: #handle multiple Nones in key
+        if k is not None:
+            key_ += (k if isinstance(k, slice) else slice(k, k + 1, None),)
+            mask += (isinstance(k, int),)
+    return key_, mask
 
 
 def get_ndarray_start_stop(ndim, key, shape):
@@ -1469,6 +1470,7 @@ def __getitem__(  # noqa: C901
                [3.3333, 3.3333, 3.3333, 3.3333, 3.3333]])
         """
         # First try some fast paths for common cases
+        newaxes = None
         if isinstance(key, np.integer):
             # Massage the key to a tuple and go the fast path
             key_ = (slice(key, key + 1), *(slice(None),) * (self.ndim - 1))
@@ -1519,6 +1521,8 @@ def __getitem__(  # noqa: C901
             start, stop, step = get_ndarray_start_stop(self.ndim, key_, self.shape)
             shape = np.array([sp - st for st, sp in zip(start, stop, strict=True)])
             shape = tuple(shape[[not m for m in mask]])
+            # Add new axes if necessary
+            newaxes = tuple(i for i, k in enumerate(key) if k is None)
 
         # Create the array to store the result
         arr = np.empty(shape, dtype=self.dtype)
@@ -1533,9 +1537,9 @@ def __getitem__(  # noqa: C901
             self._last_read.clear()
             inmutable_key = make_key_hashable(key)
             self._last_read[inmutable_key] = nparr
-
-        return nparr
-
+            
+        return nparr.expand_dims(newaxes) if newaxes is not None else nparr
+    
     def __setitem__(self, key: int | slice | Sequence[slice], value: object):
         """Set a slice of the array.
 

tests/ndarray/test_reductions.py

@@ -418,3 +418,10 @@ def test_save_constructor_reduce2(shape, disk, compute):
         blosc2.remove_urlpath(urlpath_a)
         blosc2.remove_urlpath(urlpath_b)
         blosc2.remove_urlpath("out.b2nd")
+
+def test_reduction_index():
+    shape = (20, 20)
+    a = blosc2.linspace(0, 20, num=np.prod(shape), shape=shape)
+    expr = blosc2.sum(a, axis=0)
+    assert expr[:10].shape == (10,)
+    assert expr[0].shape == ()