Merge branch 'main' of github.com:Blosc/python-blosc2

Author: FrancescAlted

src/blosc2/lazyexpr.py

@@ -3109,6 +3109,7 @@ def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=No
         # Validate the expression
         validate_expr(expression)
         expression = convert_to_slice(expression)
+        chunks, blocks = None, None
         if guess:
             # The expression has been validated, so we can evaluate it
             # in guessing mode to avoid computing reductions
@@ -3139,6 +3140,11 @@ def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=No
                 # (e.g. all operands are numpy arrays or constructors)
                 # or passed "a", "a[:10]", 'sum(a)'
                 expression_, operands_ = conserve_functions(_expression, _operands, local_vars)
+                if hasattr(new_expr, "chunks") and new_expr.chunks != (1,) * len(_shape):
+                    # for constructors with chunks in kwargs, chunks will be specified
+                    # for general expression new_expr is just with dummy scalar variables (so ignore)
+                    chunks = new_expr.chunks
+                    blocks = new_expr.blocks
             new_expr = cls(None)
             new_expr.expression = f"({expression_})"  # force parenthesis
             new_expr.operands = operands_
@@ -3147,6 +3153,8 @@ def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=No
             # Cache the dtype and shape (should be immutable)
             new_expr._dtype = _dtype
             new_expr._shape = _shape
+            if chunks is not None and blocks is not None:
+                new_expr._chunks, new_expr._blocks = chunks, blocks
         else:
             # Create a new LazyExpr object
             new_expr = cls(None)

src/blosc2/ndarray.py

@@ -5208,7 +5208,7 @@ def arange_fill(inputs, output, offset):
 
     if is_inside_new_expr() or NUM < 0:
         # We already have the dtype and shape, so return immediately
-        return blosc2.zeros(shape, dtype=dtype)
+        return blosc2.zeros(shape, dtype=dtype, **kwargs)
 
     lshape = (math.prod(shape),)
     lazyarr = blosc2.lazyudf(arange_fill, (start, stop, step), dtype=dtype, shape=lshape)
@@ -5309,7 +5309,7 @@ def linspace_fill(inputs, output, offset):
 
     if is_inside_new_expr() or num == 0:
         # We already have the dtype and shape, so return immediately
-        return blosc2.zeros(shape, dtype=dtype)  # will return empty array for num == 0
+        return blosc2.zeros(shape, dtype=dtype, **kwargs)  # will return empty array for num == 0
 
     inputs = (start, stop, num, endpoint)
     lazyarr = blosc2.lazyudf(linspace_fill, inputs, dtype=dtype, shape=(num,))

src/blosc2/shape_utils.py

@@ -521,6 +521,9 @@ def visit_Constant(self, node):
     def visit_Tuple(self, node):
         return tuple(self.visit(arg) for arg in node.elts)
 
+    def visit_List(self, node):
+        return self.visit_Tuple(node)
+
     def visit_BinOp(self, node):
         left = self.visit(node.left)
         right = self.visit(node.right)

tests/ndarray/test_lazyexpr.py

@@ -1345,7 +1345,7 @@ def test_sort():
 
 def test_listargs():
     # lazyexpr tries to convert [] to slice, but could
-    # have problems for argumetns which are lists
+    # have problems for arguments which are lists
     shape = (20,)
     na = np.arange(shape[0])
     a = blosc2.asarray(na)
@@ -1354,6 +1354,18 @@ def test_listargs():
     np.testing.assert_array_equal(expr[:], np.stack([a[:], b[:]]))
 
 
+def test_str_constructors():
+    shape = (1000, 1)
+    chunks = (100, 1)
+    a = blosc2.lazyexpr(f"linspace(0, 100, {np.prod(shape)}, shape={shape}, chunks={chunks})")
+    assert a.chunks == chunks
+    b = blosc2.lazyexpr("a.T")  # this fails unless chunkshape is assigned to a on creation
+
+    b = blosc2.ones((1000, 10))
+    a = blosc2.lazyexpr(f"b + linspace(0, 100, {np.prod(shape)}, shape={shape}, chunks={chunks})")
+    assert a.shape == np.broadcast_shapes(shape, b.shape)
+
+
 @pytest.mark.parametrize(
     "obj",
     [
@@ -1735,6 +1747,12 @@ def test_lazylinalg():
     npres = np.stack((npx, npy), axis=0)
     assert out.shape == npres.shape
     np.testing.assert_array_almost_equal(out[()], npres)
+    # --- stack ---
+    # repeat with list arg instead of tuple
+    out = blosc2.lazyexpr("stack([x, y], axis=0)")
+    npres = np.stack((npx, npy), axis=0)
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
 
     # --- tensordot ---
     out = blosc2.lazyexpr("tensordot(A, B, axes=1)")  # test with int axes