def test_binary_expression_valid_success():
json_code = {
"type": "BinaryExpression",
"operator": "*",
"left": {
"type": "Identifier",
"name": "column1"
},
"right": {
"type": "Identifier",
"name": "column2"
}
}
params = {}
expr = Expression(json_code, params)
result, msg = compare_ast(
ast.parse(expr.parsed_expression),
ast.parse("functions.col('column1') * functions.col('column2')"))
assert result, msg
json_code['operator'] = '/'
json_code['left']['type'] = 'Literal'
json_code['left']['value'] = 100
json_code['left']['raw'] = '100'
expr = Expression(json_code, params)
result, msg = compare_ast(ast.parse(expr.parsed_expression),
ast.parse("100 / functions.col('column2')"))
assert result, msg
def test_binary_expression_with_params_success():
json_code = {
"type": "BinaryExpression",
"operator": "*",
"left": {
"type": "Identifier",
"name": "column1"
},
"right": {
"type": "Identifier",
"name": "column2"
}
}
params = {'input': 'df00'}
expr = Expression(json_code, params)
result, msg = compare_ast(ast.parse(expr.parsed_expression),
ast.parse("df00['column1'] * df00['column2']"))
assert result, msg
json_code['operator'] = '/'
json_code['left']['type'] = 'Literal'
json_code['left']['value'] = 100
json_code['left']['raw'] = '100'
expr = Expression(json_code, params)
result, msg = compare_ast(ast.parse(expr.parsed_expression),
ast.parse("100 / df00['column2']"))
assert result, msg
def generate_code(self):
input_data = self.named_inputs['input data']
params = {'input': input_data}
filters = [
"({0} {1} {2})".format(f['attribute'], f['f'],
f.get('value', f.get('alias')))
for f in self.filter
]
code = """
{out} = {input}""".format(out=self.output,
input=self.named_inputs['input data'])
expressions = []
for i, expr in enumerate(self.advanced_filter):
expression = Expression(expr['tree'], params)
expressions.append(expression.parsed_expression)
if len(expressions) > 0:
for e in expressions:
code += """
{out} = {out}[{out}.apply({expr}, axis=1)]""".format(out=self.output,
expr=e)
indentation = " and "
if len(filters) > 0:
code += """
{out} = {out}.query('{f}')""".format(out=self.output,
f=indentation.join(filters))
return dedent(code)
def test_get_windows_function_success():
json_code = {
"type":
"CallExpression",
"arguments": [{
"type": "Identifier",
"name": "created_at"
}, {
"type": "Literal",
"value": 10,
"raw": "10"
}, {
"type": "Literal",
"value": "end",
"raw": "'end'"
}],
"callee": {
"type": "Identifier",
"name": "window"
}
}
params = {}
expr = Expression(json_code, params)
expected_code = ("functions.window("
"functions.col('created_at'),"
"str('10 seconds')).start.cast('timestamp')")
result, msg = compare_ast(ast.parse(expr.parsed_expression),
ast.parse(expected_code))
assert result, msg + format_code_comparison(expr.parsed_expression,
expected_code)
def test_conditional_expression_success():
json_code = {
"type": "ConditionalExpression",
"test": {
"type": "BinaryExpression",
"operator": ">",
"left": {
"type": "Identifier",
"name": "a"
},
"right": {
"type": "Literal",
"value": 1,
"raw": "1"
}
},
"consequent": {
"type": "Literal",
"value": 2,
"raw": "2"
},
"alternate": {
"type": "Literal",
"value": 3,
"raw": "3"
}
}
params = {'input': 'df00'}
expr = Expression(json_code, params)
expected_code = "functions.when((df00['a'] > 1), 2).otherwise(3)"
result, msg = compare_ast(ast.parse(expr.parsed_expression),
ast.parse(expected_code))
assert result, msg + format_code_comparison(expr.parsed_expression,
expected_code)
json_code['consequent']['type'] = 'Identifier'
json_code['consequent']['value'] = 'ok'
json_code['consequent']['name'] = 'ok'
json_code['consequent']['raw'] = '"ok"'
expr = Expression(json_code, params)
expected_code = "functions.when(df00['a'] > 1, df00['ok']).otherwise(3)"
result, msg = compare_ast(ast.parse(expr.parsed_expression),
ast.parse(expected_code))
assert result, msg + format_code_comparison(expr.parsed_expression,
expected_code)
def test_logical_expression_success():
json_code = {
"type": "LogicalExpression",
"operator": "&&",
"left": {
"type": "LogicalExpression",
"operator": "||",
"left": {
"type": "Identifier",
"name": "a"
},
"right": {
"type": "Identifier",
"name": "b"
}
},
"right": {
"type": "UnaryExpression",
"operator": "!",
"argument": {
"type": "Identifier",
"name": "c"
},
"prefix": True
}
}
params = {'input': 'df00'}
expr = Expression(json_code, params)
result, msg = compare_ast(
ast.parse(expr.parsed_expression),
ast.parse("(df00['a'] | df00['b']) & ~df00['c']"))
assert result, msg + expr.parsed_expression
json_code['operator'] = '||'
json_code['left']['type'] = 'Literal'
json_code['left']['value'] = 100
json_code['left']['raw'] = '100'
expr = Expression(json_code, params)
result, msg = compare_ast(ast.parse(expr.parsed_expression),
ast.parse("(100) | ~ (df00['c'])"))
assert result, msg + expr.parsed_expression
def test_binary_call_expression_with_params_success():
json_code = {
"type":
"CallExpression",
"arguments": [{
"type": "BinaryExpression",
"operator": "*",
"left": {
"type": "Identifier",
"name": "column1"
},
"right": {
"type": "Identifier",
"name": "column2"
}
}, {
"type": "Literal",
"value": 20,
"raw": "20"
}],
"callee": {
"type": "Identifier",
"name": "pow"
}
}
params = {'input': 'df00'}
expr = Expression(json_code, params)
result, msg = compare_ast(
ast.parse(expr.parsed_expression),
ast.parse("functions.pow(df00['column1'] * df00['column2'], 20)"))
assert result, msg
json_code['operator'] = '/'
json_code['arguments'][0]['type'] = 'Literal'
json_code['arguments'][0]['value'] = 100
json_code['arguments'][0]['raw'] = '100'
expr = Expression(json_code, params)
result, msg = compare_ast(ast.parse(expr.parsed_expression),
ast.parse("functions.pow(100, 20)"))
assert result, msg
def test_unary_expression_valid_success():
json_code = {
"type": "UnaryExpression",
"operator": "~",
"argument": {
"type": "Identifier",
"name": "column"
},
"prefix": True
}
params = {}
expr = Expression(json_code, params)
code = expr.parsed_expression
expected_code = "(~ functions.col('column'))"
result, msg = compare_ast(ast.parse(code), ast.parse(expected_code))
assert result, msg + format_code_comparison(code, expected_code)
json_code['operator'] = '-'
json_code['argument']['type'] = 'Literal'
json_code['argument']['value'] = 1
json_code['argument']['raw'] = '1'
expr = Expression(json_code, params)
result, msg = compare_ast(ast.parse(expr.parsed_expression),
ast.parse("-1"))
assert result, msg
json_code['operator'] = '+'
json_code['argument']['type'] = 'Literal'
json_code['argument']['value'] = 'some text'
json_code['argument']['raw'] = "'some text'"
expr = Expression(json_code, params)
code = expr.parsed_expression
expected_code = "+ 'some text'"
result, msg = compare_ast(ast.parse(code), ast.parse(expected_code))
assert result, msg + format_code_comparison(code, expected_code)
def test_unknown_type_expression_failure():
with pytest.raises(ValueError):
json_code = {
"type": "InvalidBinaryExpression",
"operator": ">",
"left": {
"type": "Identifier",
"name": "a"
},
"right": {
"type": "Literal",
"value": 1,
"raw": "1"
}
}
Expression(json_code, {})
def generate_code(self):
# Builds the expression and identify the target column
params = {'input': self.named_inputs['input data']}
functions = ""
for expr in self.expressions:
expression = expr['tree']
expression = Expression(expression, params)
f = expression.parsed_expression
functions += "['{}', {}],".format(expr['alias'], f)
# row.append(expression.imports) #TODO: by operation itself
code = """
{out} = {input}.copy()
functions = [{expr}]
for col, function in functions:
{out}[col] = {out}.apply(function, axis=1)
""".format(out=self.output,
input=self.named_inputs['input data'],
expr=functions)
return dedent(code)