async def test_with_nothing(hep_tables_make_local_call): # NOQA
# Hand it a xAOD that isn't based on anything.
df = DataFrame()
x = xaod_runner()
r = await x.process(df)
assert isinstance(r, DataFrame)
hep_tables_make_local_call.assert_not_called()
async def test_split_double(good_xaod, hep_tables_make_local_call): # NOQA
x = xaod_runner()
df1 = good_xaod.x + good_xaod.y
r = await x.process(df1[df1 < 20])
assert r is not None
assert isinstance(r, DataFrame)
assert hep_tables_make_local_call.call_count == 2
async def test_split_call_binary_const(good_xaod,
hep_tables_make_local_call): # NOQA
x = xaod_runner()
df1 = good_xaod.x * 10
r = await x.process(df1)
assert r is not None
assert isinstance(r, result)
assert hep_tables_make_local_call.call_count == 1
async def test_request_twice(good_xaod,
hep_tables_make_local_call_pause): # NOQA
x = xaod_runner()
ok_events = good_xaod[good_xaod.x > 10]
df1 = ok_events.x + ok_events.y + ok_events.x
_ = await x.process(df1)
assert hep_tables_make_local_call_pause.call_count == 2
async def test_split_in_filter(good_xaod, hep_tables_make_local_call): # NOQA
x = xaod_runner()
ok_events = good_xaod[good_xaod.x > good_xaod.y]
df1 = ok_events.z
r = await x.process(df1)
assert r is not None
assert isinstance(r, DataFrame)
assert hep_tables_make_local_call.call_count == 3
async def test_split_in_sqrt(good_xaod, hep_tables_make_local_call): # NOQA
x = xaod_runner()
import numpy as np
df1 = cast(DataFrame, np.sqrt(good_xaod.x + good_xaod.y)) # type: ignore
r = await x.process(df1)
assert r is not None
assert isinstance(r, DataFrame)
assert hep_tables_make_local_call.call_count == 2
assert isinstance(r.child_expr, ast.Call)
async def test_with_result(good_xaod, hep_tables_make_local_call): # NOQA
# Hand it something it can do, and then see if it can make it work!
x = xaod_runner()
df1 = good_xaod.x
r = await x.process(df1)
assert r is not None
assert isinstance(r, result)
assert r.result is not None
hep_tables_make_local_call.assert_called_once_with(df1)
async def test_split_call_binary(good_xaod,
hep_tables_make_local_call): # NOQA
x = xaod_runner()
df1 = good_xaod.x + good_xaod.y
r = await x.process(df1)
assert r is not None
assert isinstance(r, DataFrame)
assert hep_tables_make_local_call.call_count == 2
a = r.child_expr
assert isinstance(a, ast.BinOp)
assert isinstance(a.left, ast_awkward)
assert isinstance(a.right, ast_awkward)
async def test_split_call_compare(good_xaod,
hep_tables_make_local_call): # NOQA
x = xaod_runner()
df1 = good_xaod.x > good_xaod.y
r = await x.process(df1)
assert r is not None
assert isinstance(r, Column)
assert hep_tables_make_local_call.call_count == 2
a = r.child_expr
assert isinstance(a, ast.Compare)
assert isinstance(a.left, ast_awkward)
assert isinstance(a.comparators[0], ast_awkward)
async def test_run_twice(good_xaod, hep_tables_make_local_call): # NOQA
x = xaod_runner()
ok_events = good_xaod[good_xaod.x > 10]
df1 = ok_events.x + ok_events.y
_ = await x.process(df1)
r2 = await x.process(df1)
assert r2 is not None
assert isinstance(r2, DataFrame)
assert hep_tables_make_local_call.call_count == 4
a = r2.child_expr
assert isinstance(a, ast.BinOp)
assert isinstance(a.left, ast_awkward)
assert isinstance(a.right, ast_awkward)
async def test_run_mapseq(good_xaod, hep_tables_make_local_call): # NOQA
from hl_tables.atlas import a_3v
truth = good_xaod.TruthParticles('TruthParticles')
llp_truth = truth[truth.pdgId == 35]
llp_good_truth = llp_truth[llp_truth.hasProdVtx & llp_truth.hasDecayVtx]
l_prod = a_3v(llp_good_truth.prodVtx)
l_decay = a_3v(llp_good_truth.decayVtx)
lxy = (l_decay - l_prod).xy
lxy_2 = lxy[lxy.Count() == 2]
has_1muon = lxy_2[lxy_2.mapseq(lambda s: s[0] > 1 | s[1] < 2)].Count()
x = xaod_runner()
r = await x.process(has_1muon)
assert isinstance(r, DataFrame)
all_args = hep_tables_make_local_call.call_args_list
for a in all_args:
d = a[0][0]
from dataframe_expressions import render
r, _ = render(d)
print(ast.dump(r))
assert len(all_args) == 4
