Almost fixed operand handling

Author: Luke Shaw

src/blosc2/lazyexpr.py

@@ -630,49 +630,69 @@ def visit_Call(self, node):
     return set(visitor.operands)
 
 
-def conserve_functions(
-    expression: str,
-    operands_old: dict[str : blosc2.ndarray | blosc2.LazyExpr],
-    operands_new: dict[str : blosc2.ndarray | blosc2.LazyExpr],
-) -> tuple(str, dict[str : blosc2.ndarray.NDArray | blosc2.LazyExpr]):
-    operand_to_key = {id(v): k for k, v in operands_new.items()}
-
-    class OperandVisitor(ast.NodeVisitor):
-        def __init__(self):
-            self.operandset = set()
-            self.operands = {}
-            self.opcounter = 0
-            self.function_names = set()
-
-        def visit_Name(self, node):
-            if node.id == "np":
-                # Skip NumPy namespace (e.g. np.int8, which will be treated separately)
-                return
-            if node.id in self.function_names:
-                # Skip function names
-                return
-            elif (node.id not in dtype_symbols) and (node.id not in self.operandset):
-                k = operand_to_key[id(operands_old[node.id])]
-                newkey = f"o{self.opcounter}"
-                self.operands[newkey] = operands_new[k]
-                node.id = newkey
-                self.operandset.add(node.id)
-                self.opcounter += 1
-            else:
-                pass
-            self.generic_visit(node)
-
-        def visit_Call(self, node):
-            if isinstance(
-                node.func, ast.Name
-            ):  # visits Call first, then Name, so don't increment operandcounter yet
-                self.function_names.add(node.func.id)
-            self.generic_visit(node)
-
-    tree = ast.parse(expression)
-    visitor = OperandVisitor()
-    visitor.visit(tree)
-    return ast.unparse(tree), visitor.operands
+#
+# def cons_functions(
+#     expression: str,
+#     operands_old: dict[str : blosc2.ndarray | blosc2.LazyExpr],
+#     operands_new: dict[str : blosc2.ndarray | blosc2.LazyExpr],
+# ) -> tuple(str, dict[str : blosc2.ndarray.NDArray | blosc2.LazyExpr]):
+#     operand_to_key = {id(v): k for k, v in operands_new.items()}
+#     for k, v in operands_old.items():  # extend operands_to_key with old operands
+#         try:
+#             operand_to_key[id(v)]
+#         except KeyError:
+#             operand_to_key[id(v)] = k
+#             operands_new[k] = v
+#
+#     class OperandVisitor(ast.NodeVisitor):
+#         def __init__(self):
+#             self.operandmap = {}
+#             self.operands = {}
+#             self.opcounter = 0
+#             self.function_names = set()
+#
+#         def update_func(self, localop):
+#             k = operand_to_key[id(localop)]
+#             if k not in self.operandmap:
+#                 newkey = f"o{self.opcounter}"
+#                 self.operands[newkey] = operands_new[k]
+#                 self.operandmap[k] = newkey
+#                 self.opcounter += 1
+#                 return newkey
+#             else:
+#                 return self.operandmap[k]
+#
+#         def visit_Name(self, node):
+#             if node.id == "np":  # Skip NumPy namespace (e.g. np.int8, which will be treated separately)
+#                 return
+#             if node.id in self.function_names:  # Skip function names
+#                 return
+#             elif node.id not in dtype_symbols:
+#                 localop = operands_old[node.id]
+#                 if isinstance(localop, blosc2.LazyExpr):
+#                     for (
+#                         _,
+#                         v,
+#                     ) in localop.operands.items():  # expression operands already in terms of basic operands
+#                         _ = self.update_func(v)
+#                     node.id = localop.expression
+#                 else:
+#                     node.id = self.update_func(localop)
+#             else:
+#                 pass
+#             self.generic_visit(node)
+#
+#         def visit_Call(self, node):
+#             if isinstance(
+#                 node.func, ast.Name
+#             ):  # visits Call first, then Name, so don't increment operandcounter yet
+#                 self.function_names.add(node.func.id)
+#             self.generic_visit(node)
+#
+#     tree = ast.parse(expression)
+#     visitor = OperandVisitor()
+#     visitor.visit(tree)
+#     return ast.unparse(tree), visitor.operands
 
 
 class TransformNumpyCalls(ast.NodeTransformer):
@@ -2725,12 +2745,12 @@ def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=No
             if isinstance(new_expr, blosc2.LazyExpr):
                 # DO NOT restore the original expression and operands
                 # Instead rebase operands and restore only constructors
-                expression_, operands_ = conserve_functions(
-                    _expression, _operands, new_expr.operands | local_vars
-                )
-                new_expr.expression = f"({expression_})"  # force parenthesis
+                # expression_, operands_ = cons_functions(
+                #     _expression, _operands, new_expr.operands | local_vars
+                # )
+                # new_expr.expression = f"({expression_})"  # force parenthesis
                 new_expr.expression_tosave = expression
-                new_expr.operands = operands_
+                # new_expr.operands = operands_
                 new_expr.operands_tosave = operands
             else:
                 # An immediate evaluation happened (e.g. all operands are numpy arrays)

tests/ndarray/test_lazyexpr.py

@@ -1364,6 +1364,14 @@ def test_chain_expressions():
     le4_ = blosc2.lazyexpr("(le2 & le3)", {"le2": le2_, "le3": le3_})
     assert (le4_[:] == le4[:]).all()
 
+    expr1 = blosc2.lazyexpr("arange(N) + b")
+    expr2 = blosc2.lazyexpr("a * b + 1")
+    expr = blosc2.lazyexpr("expr1 - expr2")
+    expr_final = blosc2.lazyexpr("expr * expr")
+    nres = (expr * expr)[:]
+    res = expr_final.compute()
+    np.testing.assert_allclose(res[:], nres)
+
 
 # Test the chaining of multiple persistent lazy expressions
 def test_chain_persistentexpressions():