Skip to content

Commit 181fef3

Browse files
author
Luke Shaw
committed
Enable .slice method for LazyExprs
1 parent 0d74554 commit 181fef3

2 files changed

Lines changed: 24 additions & 0 deletions

File tree

src/blosc2/lazyexpr.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2718,6 +2718,9 @@ def __getitem__(self, item):
27182718
kwargs = {"_getitem": True}
27192719
return self.compute(item, **kwargs)
27202720

2721+
def slice(self, item):
2722+
return self.compute(item) # should do a slice since _getitem = False
2723+
27212724
def __str__(self):
27222725
return f"{self.expression}"
27232726

tests/ndarray/test_lazyexpr.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1066,6 +1066,27 @@ def test_eval_item(array_fixture):
10661066
np.testing.assert_allclose(res[()], nres[0:10:2])
10671067

10681068

1069+
# Test stringeval with a slice
1070+
def test_stringeval_slice(array_fixture):
1071+
a1, a2, a3, a4, na1, na2, na3, na4 = array_fixture
1072+
expr = blosc2.lazyexpr("a1 + a2 - (a3 * a4)", operands={"a1": a1, "a2": a2, "a3": a3, "a4": a4})
1073+
nres = ne_evaluate("na1 + na2 - (na3 * na4)")[:2]
1074+
res = expr.slice(slice(0, 2))
1075+
assert isinstance(res, blosc2.ndarray.NDArray)
1076+
np.testing.assert_allclose(res[:], nres)
1077+
res = expr[:2]
1078+
assert isinstance(res, np.ndarray)
1079+
np.testing.assert_allclose(res, nres)
1080+
1081+
expr1 = blosc2.lazyexpr("a1 * a2", operands={"a1": a1, "a2": a2})
1082+
expr2 = blosc2.lazyexpr("expr1[:2] + a3[:2]")
1083+
nres = ne_evaluate("(na1 * na2) + na3")[:2]
1084+
assert isinstance(expr2, blosc2.LazyExpr)
1085+
res = expr2.compute()
1086+
assert isinstance(res, blosc2.ndarray.NDArray)
1087+
np.testing.assert_allclose(res[()], nres)
1088+
1089+
10691090
# Test get_chunk method
10701091
@pytest.mark.heavy
10711092
def test_get_chunk(array_fixture):

0 commit comments

Comments
 (0)