Added []-indexing to string lazyexprs

Author: Luke Shaw

src/blosc2/lazyexpr.py

@@ -541,7 +541,8 @@ def compute_smaller_slice(larger_shape, smaller_shape, larger_slice):
 
 # Define the patterns for validation
 validation_patterns = [
-    r"[\;\[\:]",  # Flow control characters
+    # r"[\;\[\:]",  # Flow control characters
+    r"[\;]",  # Flow control characters
     r"(^|[^\w])__[\w]+__($|[^\w])",  # Dunder methods
     r"\.\b(?!real|imag|(\d*[eE]?[+-]?\d+)|(\d*[eE]?[+-]?\d+j)|\d*j\b|(sum|prod|min|max|std|mean|var|any|all|where)"
     r"\s*\([^)]*\)|[a-zA-Z_]\w*\s*\([^)]*\))",  # Attribute patterns
@@ -592,7 +593,7 @@ def validate_expr(expr: str) -> None:
         raise ValueError(f"'{expr}' is not a valid expression.")
 
     # Check for invalid characters not covered by the tokenizer
-    invalid_chars = re.compile(r"[^\w\s+\-*/%().,=<>!&|~^]")
+    invalid_chars = re.compile(r"[^\w\s+\-*/%()[].,=<>!&|~^]")
     if invalid_chars.search(skip_quotes) is not None:
         invalid_chars = invalid_chars.findall(skip_quotes)
         raise ValueError(f"Expression {expr} contains invalid characters: {invalid_chars}")
@@ -744,6 +745,38 @@ def visit_Call(self, node):
     return newexpression, newoperands
 
 
+def convert_to_slice(expression):
+    """
+    Assumes all operands are of the form o...
+    Parameters
+    ----------
+    expression: str
+
+    Returns
+    -------
+    new_expr : str
+    """
+
+    new_expr = ""
+    skip_to_char = 0
+    for i, expr_i in enumerate(expression):
+        if i < skip_to_char:
+            continue
+        if expr_i == "[":
+            k = expression[i:].find("]")  # start checking from after [
+            slice_convert = expression[i : i + k + 1]  # include [ and ]
+            slicer = eval(f"np.s_{slice_convert}")
+            slicer = (slicer,) if isinstance(slicer, slice) else slicer  # standardise to tuple
+            if any(isinstance(el, str) for el in slicer):  # handle fields
+                raise ValueError("Cannot handle fields for slicing lazy expressions.")
+            slicer = str(slicer)
+            new_expr += ".slice(" + slicer + ")"
+            skip_to_char = i + k + 1
+        else:
+            new_expr += expr_i
+    return new_expr
+
+
 class TransformNumpyCalls(ast.NodeTransformer):
     def __init__(self):
         self.replacements = {}
@@ -2780,6 +2813,7 @@ def save(self, urlpath=None, **kwargs):
     def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=None):
         # Validate the expression
         validate_expr(expression)
+        expression = convert_to_slice(expression)
         if guess:
             # The expression has been validated, so we can evaluate it
             # in guessing mode to avoid computing reductions

tests/ndarray/test_reductions.py

@@ -430,9 +430,25 @@ def test_reduction_index():
     assert arr.shape == newarr.shape
 
 
-def test_slice_in_lazy():
+def test_slice_lazy():
     shape = (20, 20)
     a = blosc2.linspace(0, 20, num=np.prod(shape), shape=shape)
     arr = blosc2.lazyexpr("anarr.slice(slice(10,15)) + 1", {"anarr": a})
     newarr = arr.compute()
     np.testing.assert_allclose(newarr[:], a.slice(slice(10, 15))[:] + 1)
+
+
+def test_slicebrackets_lazy():
+    shape = (20, 20)
+    a = blosc2.linspace(0, 20, num=np.prod(shape), shape=shape)
+    arr = blosc2.lazyexpr("anarr[10:15] + 1", {"anarr": a})
+    newarr = arr.compute()
+    np.testing.assert_allclose(newarr[:], a[10:15] + 1)
+
+    arr = blosc2.lazyexpr("anarr[10:15, 2:9] + 1", {"anarr": a})
+    newarr = arr.compute()
+    np.testing.assert_allclose(newarr[:], a[10:15, 2:9] + 1)
+
+    arr = blosc2.lazyexpr("anarr[10:15][2:9] + 1", {"anarr": a})
+    newarr = arr.compute()
+    np.testing.assert_allclose(newarr[:], a[10:15][2:9] + 1)