<gh_stars>1-10
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')

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')
# identity # multiline (             v1 +             v1         ) # negation # abs # sqrt and * # sqrt and ** # / # % # simple index calculation # conditionals # comparisons # maximum # minimum # sin / arcsin # trigonometry # mask operations # replaces mask # adds to mask # long expression # haxx # uses list # uses dict # uses string # if construct # inline comparison # unsupported unary operator # unsupported binary operator # and # does not return an array # more than one expression # dunder method # builtins # method call # attribute access # chained comparisons # unknown operand # wrong number of arguments # not a valid expression # internal numpy error (mismatching types) # code injection (serious haxx) (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,""                        ),{}                    )()                )() # make sure we have enough recursion depth for long expression