def test_VariableResolver_scope_resolve():
namespace = {
'calc': {
'a': Variable('a', rvalues=['c/b']),
'b': Variable('b', rvalues=['GEV2(b)']),
'c': Variable('c', rvalues=['b*b'])
},
'rename': {
'c': Variable('c', rvalues=['x'])
},
'raw': {
'b': Variable('b'),
'x': Variable('x')
}
}
resolver = VariableResolver(namespace)
result = resolver.resolve_scope('calc')
assert result == ([('raw', Variable('b')),
('calc', Variable('b', rvalues=['GEV2(b)'])),
('calc', Variable('c', rvalues=['b*b'])),
('calc', Variable('a', rvalues=['c/b']))], [])
assert result[0][1][1].rval == 'GEV2(raw_b)'
assert result[0][2][1].rval == 'calc_b*calc_b'
assert result[0][3][1].rval == 'calc_c/calc_b'
def test_VariableResolver_partial_common_deps_simple():
namespace = {
'sel': {
'sel0': Variable('sel0', rvalues=['flag_d0mu']),
'sel1': Variable('sel1', rvalues=['flag_mu'])
},
'calc': {
'flag_d0mu':
Variable('flag_d0mu',
rvalues=['FLAG_D0MU(mu_isMuon, k_isMuon, pi_isMuon)']),
'flag_mu':
Variable('flag_mu', rvalues=['FLAG_MU(mu_isMuon)'])
},
'raw': {
# 'mu_TRUEID': Variable('mu_TRUEID'),
'mu_isMuon': Variable('mu_isMuon'),
'k_isMuon': Variable('k_isMuon'),
}
}
resolver = VariableResolver(namespace)
result = resolver.resolve_scope('sel', ordering=['calc', 'raw'])
assert result == ([
('raw', Variable('mu_isMuon')),
('calc', Variable('flag_mu', rvalues=['FLAG_MU(mu_isMuon)'])),
('sel', Variable('sel1', rvalues=['flag_mu']))
], [Variable('sel0', rvalues=['flag_d0mu'])])
def test_VariableResolver_vars_simple():
namespace = {
'calc': {
'a': Variable('a', rvalues=['b/c']),
'b': Variable('b', rvalues=['GEV2(b)']),
},
'rename': {
'c': Variable('c', rvalues=['x'])
},
'raw': {
'b': Variable('b'),
'x': Variable('x')
}
}
resolver = VariableResolver(namespace)
result = resolver.resolve_vars_in_scope('calc',
namespace['calc'].values(),
ordering=['calc', 'rename', 'raw'])
assert result == ([('raw', Variable('b')),
('calc', Variable('b', rvalues=['GEV2(b)'])),
('raw', Variable('x')),
('rename', Variable('c', rvalues=['x'])),
('calc', Variable('a', rvalues=['b/c']))], [])
assert result[0][4][1].rval == 'calc_b/rename_c'
def test_VariableResolver_simple_fail():
namespace = {'raw': {'a': Variable('a')}}
resolver = VariableResolver(namespace)
var = Variable('a', rvalues=['b'])
assert resolver.resolve_var('keep', var) == (False, [], [])
assert resolver._resolved_names == []
def test_Variable_sub():
var = Variable('test', rvalues=['a+b'])
var.resolved = {
'a': 'scope1_a',
'b': 'scope2_b',
}
assert var.rval == 'scope1_a+scope2_b'
def test_Variable_eq():
var1 = Variable('test', rvalues=['a+b', 'c+d'])
var2 = Variable('test', rvalues=['a+b', 'c+d'])
var1.next()
assert var1 == var2
assert var1.idx == 1
assert var2.idx == 0
def test_VariableResolver_trivial():
namespace = defaultdict(dict)
resolver = VariableResolver(namespace)
var = Variable('id')
assert resolver.resolve_var('id', Variable('id')) == (True, [('id', var)],
[('id', 'id')])
assert resolver._resolved_names == []
def test_Variable_ok():
var1 = Variable('test')
var2 = Variable('test', rvalues=['a'])
assert var1.ok is True
assert var2.ok is False
var2.resolved = {'a': 'test_a'}
assert var2.ok is True
def test_VariableResolver_full_fail():
resolver = VariableResolver({
'raw': {},
'rename': {
'a': Variable('x', rvalues=['c'])
}
})
var = Variable('x', rvalues=['a+b'])
assert resolver.resolve_var('calc', var,
ordering=['rename', 'raw']) == (False, [], [])
assert var.idx == 0
assert not var.ok
def test_VariableResolver_simple():
namespace = {'raw': {'a': Variable('a')}}
resolver = VariableResolver(namespace)
var = Variable('a', rvalues=['a'])
assert resolver.resolve_var('keep', var) == (True, [('raw', Variable('a')),
('keep', var)],
[('raw', 'a'), ('keep', 'a')])
assert resolver._resolved_names == []
# Do it again (and again) and the result should be the same
assert resolver.resolve_var('keep', var) == \
resolver.resolve_var('keep', var)
assert resolver._resolved_names == []
def test_VariableResolver_selection():
resolver = VariableResolver(
{
'raw': {
'k_PT': Variable('k_PT'),
'pi_PT': Variable('pi_PT')
},
}, ['MeV'])
var = Variable(
'sel0',
rvalues=['k_PT + pi_PT > 1400.0*MeV'],
)
resolver.resolve_var('sel', var, known_names=[('raw', 'k_PT')])
assert var.rval == 'raw_k_PT + raw_pi_PT > 1400.0*MeV'
def test_VariableResolver_skip_names_simple():
resolver = VariableResolver({'raw': {}}, ['GeV'])
var = Variable('x', rvalues=['300*GeV'])
assert resolver.resolve_var('calc', var) == (True, [('calc', var)],
[('calc', 'x')])
assert var.rval == '300*GeV'
def test_Variable_simple():
var = Variable('test')
assert var.name == 'test'
assert var.type == 'nil'
assert var.deps == {'': []}
assert var.resolved == {}
assert var.idx == 0
def test_VariableResolver_multi_scope_fail():
namespace = {
'rename': {
'x': Variable('x', rvalues=['a'])
},
'raw': {
'a': Variable('a'),
'b': Variable('b')
}
}
resolver = VariableResolver(namespace)
var = Variable('a', rvalues=['b+y'])
assert resolver.resolve_var('calc', var,
ordering=['rename', 'raw']) == (False, [
('raw', Variable('b'))
], [('raw', 'b')])
def test_VariableResolver_circular():
var = Variable('x', rvalues=['GEV2(x)'])
namespace = {
'calc': {
'x': var,
},
'raw': {
'x': Variable('x'),
}
}
resolver = VariableResolver(namespace)
assert resolver.resolve_var('calc', var,
ordering=['calc', 'raw']) == (True, [
('raw', Variable('x')), ('calc', var)
], [('raw', 'x'), ('calc', 'x')])
assert var.rval == 'GEV2(raw_x)'
def test_VariableResolver_multi_scope_literal_shadow():
namespace = {
'rename': {
'x': Variable('x', rvalues=['a'])
},
'raw': {
'a': Variable('a'),
'b': Variable('b')
},
'literals': {
'b': Variable('b', literal='343')
}
}
resolver = VariableResolver(namespace)
var = Variable('a', rvalues=['x+b'])
assert resolver.resolve_var('calc',
var,
ordering=['literals', 'rename',
'raw']) == (True, [
('raw', Variable('a')),
('rename',
Variable('x', rvalues=['a'])),
('calc', var)
], [('raw', 'a'), ('rename', 'x'),
('calc', 'a')])
resolved_var = resolver.resolve_var('calc',
var,
ordering=['literals', 'rename',
'raw'])[1][-1][1]
assert resolved_var.rval == 'rename_x+343'
def test_Variable_next():
var1 = Variable('test')
var2 = Variable('test', rvalues=['a', 'b'])
assert var1.next() is False
assert var2.next() is True
assert var2.next() is False
assert var2.next() is False
def test_VariableResolver_existing_var():
resolver = VariableResolver({})
resolver._resolved_names = [('rename', 'a')]
var = Variable('a', rvalues=['x'])
assert resolver.resolve_var('rename', var) == (
True,
[],
[] # No need to resolve since it's already loaded before
)
def test_VariableResolver_multi_scope_alt_def():
namespace = {
'rename': {
'x': Variable('x', rvalues=['a'])
},
'raw': {
'a': Variable('a'),
'b': Variable('b')
}
}
resolver = VariableResolver(namespace)
var = Variable('a', rvalues=['b+y', 'x+b'])
assert resolver.resolve_var('calc', var,
ordering=['rename', 'raw']) == (True, [
('raw', Variable('a')),
('rename', Variable('x', rvalues=['a'])),
('raw', Variable('b')), ('calc', var)
], [('raw', 'a'), ('rename', 'x'),
('raw', 'b'), ('calc', 'a')])
assert var.idx == 1
assert var.rval == 'rename_x+raw_b'
def test_Variable_literal():
var = Variable('pi', literal='3.14')
assert str(var) == 'pi := 3.14'
def test_Variable_repr():
var = Variable('test', 'Double_t', ['a+b', 'a', 'b'])
assert str(var) == 'Double_t test = a+b|a|b'
def test_VariableResolver_alternative_rvalue_dep_deep():
namespace = {
'calc': {
'other_trk':
Variable(
'other_trk',
rvalues=['VEC(trk_k, trk_pi, trk_spi)', 'VEC(trk_k, trk_pi)']),
'trk_k':
Variable('trk_k', rvalues=['FAKE(k_PT)']),
'trk_pi':
Variable('trk_pi', rvalues=['FAKE(pi_PT)']),
'trk_spi':
Variable('trk_spi', rvalues=['FAKE(spi_PT)'])
},
'raw': {
'k_PT': Variable('k_PT'),
'pi_PT': Variable('pi_PT'),
# 'spi_PT': Variable('spi_PT')
}
}
resolver = VariableResolver(namespace)
result = resolver.resolve_scope('calc', ordering=['calc', 'raw'])
assert result == ([
('raw', Variable('k_PT')),
('calc', Variable('trk_k', rvalues=['FAKE(k_PT)'])),
('raw', Variable('pi_PT')),
('calc', Variable('trk_pi', rvalues=['FAKE(pi_PT)'])),
('calc',
Variable(
'other_trk',
rvalues=['VEC(trk_k, trk_pi, trk_spi)', 'VEC(trk_k, trk_pi)'])),
], [Variable('trk_spi', rvalues=['FAKE(spi_PT)'])])
def test_VariableResolver_scope_duplicated_resolution():
namespace = {
'sel': {
'sel0': Variable('sel0', rvalues=['mu_pid > 0']),
'sel1': Variable('sel1', rvalues=['test > 0'])
},
'calc': {
'mu_pid':
Variable(
'mu_pid',
rvalues=['MU_PID(mu_true_id)', 'MU_PID(mu_is_mu, mu_pid_mu)']),
'test':
Variable('test', rvalues=['TEST(mu_pid_mu, mu_is_mu)'])
},
'rename': {
'mu_true_id': Variable('mu_true_id', rvalues=['mu_TRUEID']),
'mu_is_mu': Variable('mu_is_mu', rvalues=['mu_isMuon']),
'mu_pid_mu': Variable('mu_pid_mu', rvalues=['mu_PIDmu'])
},
'raw': {
# 'mu_TRUEID': Variable('mu_TRUEID'),
'mu_isMuon': Variable('mu_isMuon'),
'mu_PIDmu': Variable('mu_PIDmu')
}
}
resolver = VariableResolver(namespace)
result = resolver.resolve_scope('sel', ordering=['calc', 'rename', 'raw'])
assert result == ([
('raw', Variable('mu_isMuon')),
('rename', Variable('mu_is_mu', rvalues=['mu_isMuon'])),
('raw', Variable('mu_PIDmu')),
('rename', Variable('mu_pid_mu', rvalues=['mu_PIDmu'])),
('calc',
Variable(
'mu_pid',
rvalues=['MU_PID(mu_true_id)', 'MU_PID(mu_is_mu, mu_pid_mu)'])),
('sel', Variable('sel0', rvalues=['mu_pid > 0'])),
('calc', Variable('test', rvalues=['TEST(mu_pid_mu, mu_is_mu)'])),
('sel', Variable('sel1', rvalues=['test > 0']))
], [])
def test_Variable_typical():
var = Variable('test', rvalues=['a+B', 'a'])
assert var.deps == {'a+B': ['a', 'B'], 'a': ['a']}
