tests
test_expressions.py
import pytest
import sys
from textwrap import dedent
import numpy as np
OPERANDS = {
'v1': np.ma.masked_array(np.arange(1, 6), dtype='float64'),
'v2': np.ma.masked_array(2 * np.arange(1, 6), dtype='float64', mask=np.array([1, 1, 1, 0, 0])),
}
VALID_EXPR = (
# identity
('v1', OPERANDS['v1']),
# multiline
(
dedent('''
(
v1 +
v1
)
'''),
2 * OPERANDS['v1']
),
# negation
('-v1', -OPERANDS['v1']),
# abs
('abs(v1)', np.abs(OPERANDS['v1'])),
# sqrt and *
('sqrt(v1 * v1)', OPERANDS['v1']),
# sqrt and **
('sqrt(v1 ** 2)', OPERANDS['v1']),
# /
('v1 / v1', np.ones_like(OPERANDS['v1'])),
# %
('v1 % v1', np.zeros_like(OPERANDS['v1'])),
# simple index calculation
(
'(v2 - v1) / (v2 + v1)',
(OPERANDS['v2'] - OPERANDS['v1']) / (OPERANDS['v2'] + OPERANDS['v1'])
),
# conditionals
(
'where(v2 > v1, 100, 0)', np.where(OPERANDS['v2'] > OPERANDS['v1'], 100, 0)
),
# comparisons
('v1 == v2', OPERANDS['v1'] == OPERANDS['v2']),
('v1 != v2', OPERANDS['v1'] != OPERANDS['v2']),
('v1 > v2', OPERANDS['v1'] > OPERANDS['v2']),
('v1 < v2', OPERANDS['v1'] < OPERANDS['v2']),
('v1 >= v2', OPERANDS['v1'] >= OPERANDS['v2']),
('v1 < = v2', OPERANDS['v1'] < = OPERANDS['v2']),
('(v1 < 0.5) & (v2 > 0.5)', (OPERANDS['v1'] < 0.5) & (OPERANDS['v2'] > 0.5)),
('(v1 < 0.5) | (v2 > 0.5)', (OPERANDS['v1'] < 0.5) | (OPERANDS['v2'] > 0.5)),
('~(v1 < 0.5) & (v2 > 0.5)', ~(OPERANDS['v1'] < 0.5) & (OPERANDS['v2'] > 0.5)),
# maximum
(
'maximum(v1, v2)', np.maximum(OPERANDS['v1'], OPERANDS['v2'])
),
# minimum
(
'minimum(v1, v2)', np.minimum(OPERANDS['v1'], OPERANDS['v2'])
),
# sin / arcsin
(
'arcsin(sin(v1))', np.arcsin(np.sin(OPERANDS['v1']))
),
# trigonometry
(
'sin(pi * v1)', np.sin(np.pi * OPERANDS['v1'])
),
# mask operations
(
'setmask(v1, getmask(v2))', np.ma.masked_array(OPERANDS['v1'], mask=OPERANDS['v2'].mask)
),
(
'setmask(v2, nomask)', np.ma.masked_array(OPERANDS['v2'], mask=np.ma.nomask)
),
( # replaces mask
'setmask(v2, ~getmask(v2))', np.ma.masked_array(OPERANDS['v2'], mask=~OPERANDS['v2'].mask)
),
( # adds to mask
'masked_where(~getmask(v2), v2)', np.ma.masked_array(OPERANDS['v2'], mask=True)
),
# long expression
(
'+'.join(['v1'] * 1000), sum(OPERANDS['v1'] for _ in range(1000))
)
)
INVALID_EXPR = (
# haxx
('__builtins__["dir"]', 'not allowed in expressions'),
# uses list
('[0] * 1000000000', 'not allowed in expressions'),
# uses dict
('{}', 'not allowed in expressions'),
# uses string
('"general kenobi!"', 'not allowed in expressions'),
# if construct
('if True: v1', 'is not a valid expression'),
# inline comparison
('v1 if True else 0', 'not allowed in expressions'),
# unsupported unary operator
('not v1', 'unary operator Not'),
# unsupported binary operator
('v1 ^ v1', 'binary operator BitXor'),
# and
('v1 and v2', 'not allowed in expressions'),
# does not return an array
('0', 'does not return an array'),
# more than one expression
('v1; v1', 'is not a valid expression'),
# dunder method
('__name__', 'unrecognized name \'__name__\''),
# builtins
('dir(v1)', 'unrecognized name \'dir\''),
# method call
('v1.mean()', 'not allowed in expressions'),
# attribute access
('v1.size', 'not allowed in expressions'),
# chained comparisons
('where(v1 < v2 == v2, 0, 1)', 'not supported'),
# unknown operand
('v100', 'unrecognized name \'v100\''),
# wrong number of arguments
('maximum(v1, v1, v1, v1, v1)', 'got 5, expected 2'),
# not a valid expression
('k = v1', 'is not a valid expression'),
# internal numpy error (mismatching types)
('v1 & v2', 'unexpected error'),
# code injection (serious haxx)
(
dedent('''
(lambda fc=(
lambda n: [
c for c in
().__class__.__bases__[0].__subclasses__()
if c.__name__ == n
][0]
):
fc("function")(
fc("code")(
0,0,0,0,"KABOOM",(),(),(),"","",0,""
),{}
)()
)()
'''),
'not allowed in expressions'
)
)
@pytest.mark.parametrize('case', VALID_EXPR)
def test_valid_expression(case):
from terracotta.expressions import evaluate_expression
# make sure we have enough recursion depth for long expression
sys.setrecursionlimit(10_000)
expr, result = case
np.testing.assert_array_equal(
evaluate_expression(expr, OPERANDS),
result
)
@pytest.mark.parametrize('case', INVALID_EXPR)
def test_invalid_expression(case):
from terracotta.expressions import evaluate_expression
expr, exc_msg = case
with pytest.raises(ValueError) as raised_exc:
evaluate_expression(expr, OPERANDS)
assert exc_msg in str(raised_exc.value)
def test_invalid_compop(monkeypatch):
from terracotta.expressions import evaluate_expression, ExpressionParser
expr = 'v0 < v1'
exc_msg = 'comparison operator'
with monkeypatch.context() as m:
m.setattr(ExpressionParser, 'NODE_TO_COMPOP', {})
with pytest.raises(ValueError) as raised_exc:
evaluate_expression(expr, OPERANDS)
assert exc_msg in str(raised_exc.value)
def test_timeout():
from terracotta.expressions import evaluate_expression
with pytest.raises(RuntimeError) as raised_exc:
evaluate_expression('+'.join(['v1'] * 100), {'v1': np.ones((256, 256))}, timeout=0)
assert 'timeout' in str(raised_exc.value)
def test_mask_invalid():
from terracotta.expressions import evaluate_expression
res = evaluate_expression('where(v1 + v2 < 10, nan, 0)', OPERANDS)
mask = (OPERANDS['v1'] + OPERANDS['v2'] < 10) | OPERANDS['v1'].mask | OPERANDS['v2'].mask
assert isinstance(res, np.ma.MaskedArray)
assert np.all(res == 0)
assert np.array_equal(res.mask, mask)
def test_out_dtype():
from terracotta.expressions import evaluate_expression
operands = dict(v1=np.ones(10, dtype='int64'), v2=np.zeros(10, dtype='int32'))
res = evaluate_expression('v1 + v2', operands)
assert isinstance(res, np.ma.MaskedArray)
assert res.dtype == np.dtype('int64')
