Skip to content

Commit

Permalink
Merge pull request #127 from tdixon97/main
Browse files Browse the repository at this point in the history
  • Loading branch information
gipert authored Jan 3, 2025
2 parents 4e548e8 + f5ed430 commit f1a5297
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 2 deletions.
17 changes: 15 additions & 2 deletions src/lgdo/types/table.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

import logging
from collections.abc import Mapping
from types import ModuleType
from typing import Any
from warnings import warn

Expand Down Expand Up @@ -266,6 +267,7 @@ def eval(
self,
expr: str,
parameters: Mapping[str, str] | None = None,
modules: Mapping[str, ModuleType] | None = None,
) -> LGDO:
"""Apply column operations to the table and return a new LGDO.
Expand Down Expand Up @@ -299,6 +301,10 @@ def eval(
a dictionary of function parameters. Passed to
:func:`numexpr.evaluate`` as `local_dict` argument or to
:func:`eval` as `locals` argument.
modules
a dictionary of additional modules used by the expression. If this is not `None`
then :func:`eval`is used and the expression can depend on any modules from this dictionary in
addition to awkward and numpy. These are passed to :func:`eval` as `globals` argument.
Examples
--------
Expand Down Expand Up @@ -339,8 +345,8 @@ def eval(
msg = f"evaluating {expr!r} with locals={(self_unwrap | parameters)} and {has_ak=}"
log.debug(msg)

# use numexpr if we are only dealing with numpy data types
if not has_ak:
# use numexpr if we are only dealing with numpy data types (and no global dictionary)
if not has_ak and modules is None:
out_data = ne.evaluate(
expr,
local_dict=(self_unwrap | parameters),
Expand All @@ -366,6 +372,9 @@ def eval(

# resort to good ol' eval()
globs = {"ak": ak, "np": np}
if modules is not None:
globs = globs | modules

out_data = eval(expr, globs, (self_unwrap | parameters))

msg = f"...the result is {out_data!r}"
Expand All @@ -380,6 +389,10 @@ def eval(
if np.isscalar(out_data):
return Scalar(out_data)

# if out_data is already an LGDO just return it
if isinstance(out_data, LGDO):
return out_data

msg = (
f"evaluation resulted in a {type(out_data)} object, "
"I don't know which LGDO this corresponds to"
Expand Down
11 changes: 11 additions & 0 deletions tests/types/test_table_eval.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
from __future__ import annotations

import hist
import numpy as np
import pytest

import lgdo

Expand Down Expand Up @@ -31,6 +33,7 @@ def test_eval_dependency():
),
}
)

r = obj.eval("sum(a)")
assert isinstance(r, lgdo.Scalar)

Expand Down Expand Up @@ -77,3 +80,11 @@ def test_eval_dependency():
assert isinstance(r, lgdo.Scalar)

assert obj.eval("np.sum(a) + ak.sum(e)")

# test with modules argument, the simplest is using directly lgdo
res = obj.eval("lgdo.Array([1,2,3])", {}, modules={"lgdo": lgdo})
assert res == lgdo.Array([1, 2, 3])

# check bad type
with pytest.raises(RuntimeError):
obj.eval("hist.Hist()", modules={"hist": hist})

0 comments on commit f1a5297

Please sign in to comment.