Merge pull request #402 from Blosc/fixIndexing

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

Author: FrancescAlted

src/blosc2/lazyexpr.py

@@ -2440,16 +2440,17 @@ 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 self._output
+                return lazy_expr[item]
 
             return chunked_eval(lazy_expr.expression, lazy_expr.operands, item, **kwargs)
 

tests/ndarray/test_reductions.py

@@ -418,3 +418,11 @@ 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 == ()