diff --git a/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py index 4fdf781f..0ce8d978 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -16,6 +16,7 @@ import os import pathlib import re +import sys import threading from abc import ABC, abstractmethod from enum import Enum @@ -2209,11 +2210,48 @@ def lazyudf( return LazyUDF(func, inputs, dtype, chunked_eval, **kwargs) +def seek_operands(names, local_dict=None, global_dict=None, _frame_depth: int = 2): + """ + Get the arguments based on the names. + """ + call_frame = sys._getframe(_frame_depth) + + clear_local_dict = False + if local_dict is None: + local_dict = call_frame.f_locals + clear_local_dict = True + try: + frame_globals = call_frame.f_globals + if global_dict is None: + global_dict = frame_globals + + # If `call_frame` is the top frame of the interpreter we can't clear its + # `local_dict`, because it is actually the `global_dict`. + clear_local_dict = clear_local_dict and frame_globals is not local_dict + + op_dict = {} + for name in names: + try: + a = local_dict[name] + except KeyError: + a = global_dict[name] + op_dict[name] = a + finally: + # If we generated local_dict via an explicit reference to f_locals, + # clear the dict to prevent creating extra ref counts in the caller's scope + if clear_local_dict and hasattr(local_dict, "clear"): + local_dict.clear() + + return op_dict + + def lazyexpr( expression: str | bytes | LazyExpr, operands: dict | None = None, out: blosc2.NDArray | np.ndarray = None, where: tuple | list | None = None, + local_dict: dict | None = None, + global_dict: dict | None = None, ) -> LazyExpr: """ Get a LazyExpr from an expression. @@ -2227,17 +2265,18 @@ def lazyexpr( operands: dict The dictionary with operands. Supported values are NumPy.ndarray, Python scalars, :ref:`NDArray`, :ref:`NDField` or :ref:`C2Array` instances. + If None, the operands will be seeked in the local and global dictionaries. out: NDArray or np.ndarray, optional The output array where the result will be stored. If not provided, a new array will be created. where: tuple, list, optional A sequence of arguments for the where clause in the expression. - guess: bool, optional - Whether to guess the output dtype and shape. If False, the dtype and shape - will be computed producing temporary arrays in the process (e.g. for reductions). - If True, the dtype and shape will be guessed from the expression, but without - evaluating any part of it. Use True when you want to e.g. save the expression - but without evaluating it. + local_dict: dict, optional + The local dictionary to use when looking for operands in the expression. + If not provided, the local dictionary of the caller will be used. + global_dict: dict, optional + The global dictionary to use when looking for operands in the expression. + If not provided, the global dictionary of the caller will be used. Returns ------- @@ -2279,7 +2318,13 @@ def lazyexpr( expression._where_args = where_args return expression if operands is None: - raise ValueError("`operands` must be provided for a string expression") + # Try to get operands from variables in the stack + operands = get_expr_operands(expression) + # If no operands are found, raise an error + if operands is None: + raise ValueError("No operands found in the expression") + # Look for operands in the stack + operands = seek_operands(operands, local_dict, global_dict) return LazyExpr._new_expr(expression, operands, guess=True, out=out, where=where) diff --git a/tests/ndarray/test_lazyexpr.py b/tests/ndarray/test_lazyexpr.py index d4db5d79..5dd16da0 100644 --- a/tests/ndarray/test_lazyexpr.py +++ b/tests/ndarray/test_lazyexpr.py @@ -783,7 +783,8 @@ def test_broadcasting(broadcast_fixture): ("numpy", "numpy"), ], ) -def test_lazyexpr(array_fixture, operand_mix): +@pytest.mark.parametrize("operand_guess", [True, False]) +def test_lazyexpr(array_fixture, operand_mix, operand_guess): a1, a2, a3, a4, na1, na2, na3, na4 = array_fixture if operand_mix[0] == "NDArray" and operand_mix[1] == "NDArray": operands = {"a1": a1, "a2": a2, "a3": a3, "a4": a4} @@ -795,7 +796,10 @@ def test_lazyexpr(array_fixture, operand_mix): operands = {"a1": na1, "a2": na2, "a3": na3, "a4": na4} # Check eval() - expr = blosc2.lazyexpr("a1 + a2 - a3 * a4", operands=operands) + if operand_guess: + expr = blosc2.lazyexpr("a1 + a2 - a3 * a4") + else: + expr = blosc2.lazyexpr("a1 + a2 - a3 * a4", operands=operands) nres = ne.evaluate("na1 + na2 - na3 * na4") res = expr.compute() np.testing.assert_allclose(res[:], nres)