Improve shape-parser for various args

Author: lshaw8317

src/blosc2/shape_utils.py

@@ -391,14 +391,7 @@ def visit_Call(self, node):  # noqa : C901
         # --- Parse keyword args ---
         kwargs = {}
         for kw in node.keywords:
-            if isinstance(kw.value, ast.Constant):
-                kwargs[kw.arg] = kw.value.value
-            elif isinstance(kw.value, ast.Tuple):
-                kwargs[kw.arg] = tuple(
-                    e.value if isinstance(e, ast.Constant) else self._lookup_value(e) for e in kw.value.elts
-                )
-            else:
-                kwargs[kw.arg] = self._lookup_value(kw.value)
+            kwargs[kw.arg] = self._lookup_value(kw.value)
 
         # ------- handle linear algebra ---------------
         if base_name in linalg_funcs:
@@ -539,17 +532,62 @@ def _eval_slice(self, node):
         else:
             raise ValueError(f"Unsupported slice expression: {ast.dump(node)}")
 
-    def _lookup_value(self, node):
+    def _lookup_value(self, node):  # noqa : C901
         """Look up a value in self.shapes if node is a variable name, else constant value."""
+        # Name -> lookup in shapes mapping
         if isinstance(node, ast.Name):
             return self.shapes.get(node.id, None)
-        elif isinstance(node, ast.Constant):
+
+        # Constant -> return its value
+        if isinstance(node, ast.Constant):
             return node.value
-        elif isinstance(node, ast.Tuple):
-            return tuple(e.value for e in node.elts)
-        else:
+
+        # Tuple of constants / expressions
+        if isinstance(node, ast.Tuple):
+            vals = []
+            for e in node.elts:
+                v = self._lookup_value(e)
+                vals.append(v)
+            return tuple(vals)
+
+        # Unary operations (e.g. -1)
+        if isinstance(node, ast.UnaryOp):
+            # handle negative constants like -1
+            if isinstance(node.op, ast.USub):
+                val = self._lookup_value(node.operand)
+                if isinstance(val, (int, float)):
+                    return -val
+            # handle + (USub) if needed
+            if isinstance(node.op, ast.UAdd):
+                return self._lookup_value(node.operand)
             return None
 
+        # Simple binary ops with constant operands (e.g. 1+2)
+        if isinstance(node, ast.BinOp):
+            left = self._lookup_value(node.left)
+            right = self._lookup_value(node.right)
+            if left is None or right is None:
+                return None
+            try:
+                if isinstance(node.op, ast.Add):
+                    return left + right
+                if isinstance(node.op, ast.Sub):
+                    return left - right
+                if isinstance(node.op, ast.Mult):
+                    return left * right
+                if isinstance(node.op, ast.FloorDiv):
+                    return left // right
+                if isinstance(node.op, ast.Div):
+                    return left / right
+                if isinstance(node.op, ast.Mod):
+                    return left % right
+            except Exception:
+                return None
+            return None
+
+        # fallback
+        return None
+
 
 # --- Public API ---
 def infer_shape(expr, shapes):

tests/ndarray/test_lazyexpr.py

@@ -1702,7 +1702,9 @@ def test_lazylinalg():
     npres = np.squeeze(npD, -1)
     assert out.shape == npres.shape
     np.testing.assert_array_almost_equal(out[()], npres)
-
+    # refresh D since squeeze is in-place
+    s = shapes["D"]
+    D = blosc2.linspace(0, np.prod(s), shape=s)
     out = blosc2.lazyexpr("D.squeeze(axis=-1)")
     npres = np.squeeze(npD, -1)
     assert out.shape == npres.shape

tests/ndarray/test_squeeze.py

@@ -19,7 +19,6 @@
         ((23, 1, 1, 34), (20, 1, 1, 20), None, 1234, 2),
         ((80, 1, 51, 60, 1), None, (6, 1, 6, 26, 1), 3.333, 4),
         ((1, 1, 1), None, None, True, (1, 2)),
-        ((1, 1, 1), None, None, True, None),
     ],
 )
 def test_squeeze(shape, chunks, blocks, fill_value, axis):