Merge pull request #417 from Blosc/enableSliceLazy

Enable slice lazy

Author: FrancescAlted

src/blosc2/lazyexpr.py

@@ -2718,6 +2718,9 @@ def __getitem__(self, item):
         kwargs = {"_getitem": True}
         return self.compute(item, **kwargs)
 
+    def slice(self, item):
+        return self.compute(item)  # should do a slice since _getitem = False
+
     def __str__(self):
         return f"{self.expression}"
 

tests/ndarray/test_lazyexpr.py

@@ -1066,6 +1066,28 @@ def test_eval_item(array_fixture):
     np.testing.assert_allclose(res[()], nres[0:10:2])
 
 
+# Test lazyexpr's slice method
+def test_eval_slice(array_fixture):
+    a1, a2, a3, a4, na1, na2, na3, na4 = array_fixture
+    expr = blosc2.lazyexpr("a1 + a2 - (a3 * a4)", operands={"a1": a1, "a2": a2, "a3": a3, "a4": a4})
+    nres = ne_evaluate("na1 + na2 - (na3 * na4)")[:2]
+    res = expr.slice(slice(0, 2))
+    assert isinstance(res, blosc2.ndarray.NDArray)
+    np.testing.assert_allclose(res[:], nres)
+    res = expr[:2]
+    assert isinstance(res, np.ndarray)
+    np.testing.assert_allclose(res, nres)
+
+    # string lazy expressions automatically use .slice internally
+    expr1 = blosc2.lazyexpr("a1 * a2", operands={"a1": a1, "a2": a2})
+    expr2 = blosc2.lazyexpr("expr1[:2] + a3[:2]")
+    nres = ne_evaluate("(na1 * na2) + na3")[:2]
+    assert isinstance(expr2, blosc2.LazyExpr)
+    res = expr2.compute()
+    assert isinstance(res, blosc2.ndarray.NDArray)
+    np.testing.assert_allclose(res[()], nres)
+
+
 # Test get_chunk method
 @pytest.mark.heavy
 def test_get_chunk(array_fixture):