Repository: QuantStack/py2vega

Issue:
Assignment support
Assignments could be potentially supported.

_e.g._:
```Python
def foo(value):
    a = 36
    return a if value < 4 else 32
```

Could be transpiled to:
`"value < 4 ? 36 : 32"`

Discussion:

Patch:
diff --git a/README.md b/README.md
index ce1c037..565e998 100644
--- a/README.md
+++ b/README.md
@@ -56,3 +56,17 @@ def foo(value):
 
 foo_expr = py2vega(foo, whitelist=['value'])  # "if(isNaN(value), 'It is NaN...', value)"
 ```
+
+Even if assignments are prohibited in Vega-expressions, you can assign variables in your Python function, it will be turned into a valid Vega-expression anyway:
+
+```Python
+from py2vega import py2vega
+
+def foo(value):
+    a = 'green'
+    b = 'red'
+
+    return a if value < 3 else b
+
+foo_expr = py2vega(foo, whitelist=['value'])  # "value < 3 ? 'green' : 'red'"
+```
diff --git a/py2vega/main.py b/py2vega/main.py
index 898fa0c..3a2157e 100644
--- a/py2vega/main.py
+++ b/py2vega/main.py
@@ -8,12 +8,12 @@ from .math import math_functions
 from .constants import constants
 
 
-def return_stmt(stmt, whitelist):
+def return_stmt(stmt, whitelist, scope):
     """Turn a Python return statement into a vega-expression."""
-    return pystmt2vega(stmt.value, whitelist)
+    return pystmt2vega(stmt.value, whitelist, scope)
 
 
-def if_stmt(stmt, whitelist):
+def if_stmt(stmt, whitelist, scope):
     """Turn a Python if statement into a vega-expression."""
     def check_sanity(stmt_body):
         if len(stmt_body) > 1 and not isinstance(stmt_body[0], ast.Return) and not isinstance(stmt_body[0], ast.If):
@@ -23,13 +23,13 @@ def if_stmt(stmt, whitelist):
     check_sanity(stmt.orelse)
 
     return 'if({}, {}, {})'.format(
-        pystmt2vega(stmt.test, whitelist),
-        pystmt2vega(stmt.body[0], whitelist),
-        pystmt2vega(stmt.orelse[0], whitelist)
+        pystmt2vega(stmt.test, whitelist, scope),
+        pystmt2vega(stmt.body[0], whitelist, scope),
+        pystmt2vega(stmt.orelse[0], whitelist, scope)
     )
 
 
-def nameconstant_expr(expr, _):
+def nameconstant_expr(expr, whitelist, scope):
     """Turn a Python nameconstant expression into a vega-expression."""
     if expr.value is False:
         return 'false'
@@ -40,53 +40,67 @@ def nameconstant_expr(expr, _):
     raise NameError('name \'{}\' is not defined, only a subset of Python is supported'.format(str(expr.value)))
 
 
-def num_expr(expr, _):
+def num_expr(expr, whitelist, scope):
     """Turn a Python num expression into a vega-expression."""
     return repr(expr.n)
 
 
-def str_expr(expr, _):
+def str_expr(expr, whitelist, scope):
     """Turn a Python str expression into a vega-expression."""
     return repr(expr.s)
 
 
-def list_expr(expr, whitelist):
+def list_expr(expr, whitelist, scope):
     """Turn a Python list expression into a vega-expression."""
-    return '[{}]'.format(', '.join(pystmt2vega(elt, whitelist) for elt in expr.elts))
+    return '[{}]'.format(', '.join(pystmt2vega(elt, whitelist, scope) for elt in expr.elts))
 
 
-def dict_expr(expr, whitelist):
+def dict_expr(expr, whitelist, scope):
     """Turn a Python dict expression into a vega-expression."""
     return '{{{}}}'.format(
         ', '.join([
-            '{}: {}'.format(pystmt2vega(expr.keys[idx], whitelist), pystmt2vega(expr.values[idx], whitelist))
+            '{}: {}'.format(pystmt2vega(expr.keys[idx], whitelist, scope), pystmt2vega(expr.values[idx], whitelist, scope))
             for idx in range(len(expr.keys))
         ])
     )
 
 
-def unaryop_expr(expr, whitelist):
+def assign_expr(expr, whitelist, scope):
+    """Turn a Python assignment expression into a vega-expression. And save the assigned variable in the current scope."""
+    value = pystmt2vega(expr.value, whitelist, scope)
+
+    for target in expr.targets:
+        if not isinstance(target, ast.Name):
+            raise RuntimeError('Unsupported target {} for the assignment'.format(str(target)))
+
+        scope[target.id] = value
+
+    # Assignment in Python returns None
+    return 'null'
+
+
+def unaryop_expr(expr, whitelist, scope):
     """Turn a Python unaryop expression into a vega-expression."""
     if isinstance(expr.op, ast.Not):
-        return '!({})'.format(pystmt2vega(expr.operand, whitelist))
+        return '!({})'.format(pystmt2vega(expr.operand, whitelist, scope))
     if isinstance(expr.op, ast.USub):
-        return '-{}'.format(pystmt2vega(expr.operand, whitelist))
+        return '-{}'.format(pystmt2vega(expr.operand, whitelist, scope))
     if isinstance(expr.op, ast.UAdd):
-        return '+{}'.format(pystmt2vega(expr.operand, whitelist))
+        return '+{}'.format(pystmt2vega(expr.operand, whitelist, scope))
 
     raise RuntimeError('Unsupported {} operator, only a subset of Python is supported'.format(str(expr.op)))
 
 
-def boolop_expr(expr, whitelist):
+def boolop_expr(expr, whitelist, scope):
     """Turn a Python boolop expression into a vega-expression."""
     return '{} {} {}'.format(
-        pystmt2vega(expr.values[0], whitelist),
+        pystmt2vega(expr.values[0], whitelist, scope),
         '||' if isinstance(expr.op, ast.Or) else '&&',
-        pystmt2vega(expr.values[1], whitelist)
+        pystmt2vega(expr.values[1], whitelist, scope)
     )
 
 
-def _binop_expr_impl(left_expr, op, right_expr, whitelist=[]):
+def _binop_expr_impl(left_expr, op, right_expr, whitelist, scope):
     operator_mapping = {
         ast.Eq: '==', ast.NotEq: '!=',
         ast.Lt: '<', ast.LtE: '<=',
@@ -97,8 +111,8 @@ def _binop_expr_impl(left_expr, op, right_expr, whitelist=[]):
         ast.Mod: '%'
     }
 
-    left = left_expr if isinstance(left_expr, str) else pystmt2vega(left_expr, whitelist)
-    right = pystmt2vega(right_expr, whitelist)
+    left = left_expr if isinstance(left_expr, str) else pystmt2vega(left_expr, whitelist, scope)
+    right = pystmt2vega(right_expr, whitelist, scope)
 
     if isinstance(op, ast.In):
         return 'indexof({}, {}) != -1'.format(right, left)
@@ -115,38 +129,43 @@ def _binop_expr_impl(left_expr, op, right_expr, whitelist=[]):
     return '{} {} {}'.format(left, operator, right)
 
 
-def binop_expr(expr, whitelist):
+def binop_expr(expr, whitelist, scope):
     """Turn a Python binop expression into a vega-expression."""
-    return _binop_expr_impl(expr.left, expr.op, expr.right, whitelist)
+    return _binop_expr_impl(expr.left, expr.op, expr.right, whitelist, scope)
 
 
-def if_expr(expr, whitelist):
+def if_expr(expr, whitelist, scope):
     """Turn a Python if expression into a vega-expression."""
     return '{} ? {} : {}'.format(
-        pystmt2vega(expr.test, whitelist),
-        pystmt2vega(expr.body, whitelist),
-        pystmt2vega(expr.orelse, whitelist)
+        pystmt2vega(expr.test, whitelist, scope),
+        pystmt2vega(expr.body, whitelist, scope),
+        pystmt2vega(expr.orelse, whitelist, scope)
     )
 
 
-def compare_expr(expr, whitelist):
+def compare_expr(expr, whitelist, scope):
     """Turn a Python compare expression into a vega-expression."""
     left_operand = expr.left
 
     for idx in range(len(expr.comparators)):
-        left_operand = _binop_expr_impl(left_operand, expr.ops[idx], expr.comparators[idx], whitelist)
+        left_operand = _binop_expr_impl(left_operand, expr.ops[idx], expr.comparators[idx], whitelist, scope)
 
     return left_operand
 
 
-def name_expr(expr, whitelist):
+def name_expr(expr, whitelist, scope):
     """Turn a Python name expression into a vega-expression."""
+    # If it's in the scope, return it's evaluated expression
+    if expr.id in scope:
+        return scope[expr.id]
+
     if expr.id in constants or expr.id in whitelist:
         return expr.id
+
     raise NameError('name \'{}\' is not defined, only a subset of Python is supported'.format(expr.id))
 
 
-def call_expr(expr, whitelist):
+def call_expr(expr, whitelist, scope):
     """Turn a Python call expression into a vega-expression."""
     if isinstance(expr.func, ast.Name):
         func_name = expr.func.id
@@ -157,13 +176,13 @@ def call_expr(expr, whitelist):
     if func_name in math_functions:
         return '{}({})'.format(
             func_name,
-            ', '.join([pystmt2vega(arg, whitelist) for arg in expr.args])
+            ', '.join([pystmt2vega(arg, whitelist, scope) for arg in expr.args])
         )
 
     raise NameError('name \'{}\' is not defined, only a subset of Python is supported'.format(func_name))
 
 
-def attribute_expr(expr, _):
+def attribute_expr(expr, whitelist, scope):
     """Turn a Python attribute expression into a vega-expression."""
     return expr.attr
 
@@ -177,6 +196,7 @@ stmt_mapping = {
     ast.Tuple: list_expr,
     ast.List: list_expr,
     ast.Dict: dict_expr,
+    ast.Assign: assign_expr,
     ast.UnaryOp: unaryop_expr,
     ast.BoolOp: boolop_expr,
     ast.BinOp: binop_expr,
@@ -188,14 +208,14 @@ stmt_mapping = {
 }
 
 
-def pystmt2vega(stmt, whitelist=[]):
+def pystmt2vega(stmt, whitelist=[], scope={}):
     """Turn a Python statement object into a Vega expression."""
     func = stmt_mapping.get(stmt.__class__)
 
     if func is None:
         raise RuntimeError('Unsupported {} statement'.format(repr(stmt)))
 
-    return func(stmt, whitelist)
+    return func(stmt, whitelist, scope)
 
 
 def py2vega(value, whitelist=[]):
@@ -213,11 +233,9 @@ def py2vega(value, whitelist=[]):
 
         func = ast.parse(value, '<string>', 'exec').body[0]
 
-        if len(func.body) > 1:
-            raise RuntimeError("""
-                The only statement in your function must be a `Return` statement or an `if` statement,
-                but a value of:\n {} was given""".format(value))
-
-        return pystmt2vega(func.body[0], whitelist)
+        scope = {}
+        for stmt in func.body[:-1]:
+            pystmt2vega(stmt, whitelist, scope)
+        return pystmt2vega(func.body[-1], whitelist, scope)
 
     raise RuntimeError('py2vega only supports code string or functions as input')


Test patch:
diff --git a/test/test_py2vega.py b/test/test_py2vega.py
index 3d7ef6b..e93ab7a 100644
--- a/test/test_py2vega.py
+++ b/test/test_py2vega.py
@@ -120,31 +120,12 @@ def invalid_func1():
     print(3)
 
 
-def invalid_func2():
-    f = 3
-
-
-def invalid_func3(value):
-    value < 3
-    return 3 < value <= 4
-
-
 def test_invalid1():
     with pytest.raises(RuntimeError):
         py2vega(invalid_func1)
 
 
 def test_invalid2():
-    with pytest.raises(RuntimeError):
-        assert py2vega(invalid_func2)
-
-
-def test_invalid3():
-    with pytest.raises(RuntimeError):
-        py2vega(invalid_func3)
-
-
-def test_invalid4():
     with pytest.raises(RuntimeError):
         py2vega(lambda value: value)
 
@@ -160,3 +141,68 @@ def conditional_func(value):
 
 def test_if_stmt():
     assert py2vega(conditional_func, whitelist) == "if(value < 3, 'red', if(value < 5, 'green', 'yellow'))"
+
+
+def assign_func1(value):
+    val = ('USA', 'Japan')
+
+    return 'red' if value in val else 'green'
+
+
+def assign_func2(value):
+    a = 'green'
+    b = 'red'
+
+    return a if value < 3 else b
+
+
+def assign_func3(value):
+    a = 'green'
+    a = 'red'
+
+    return a
+
+
+def assign_func4(value):
+    a = 'green'
+    b = a
+
+    return b
+
+
+def assign_func5(value):
+    a = b = 'Hello'
+
+    return (a, b)
+
+
+def assign_func6(value):
+    a = 'Hello'
+    b = a
+    a = 'World'
+
+    return b
+
+
+def test_assign1():
+    assert py2vega(assign_func1, whitelist) == "indexof(['USA', 'Japan'], value) != -1 ? 'red' : 'green'"
+
+
+def test_assign2():
+    assert py2vega(assign_func2, whitelist) == "value < 3 ? 'green' : 'red'"
+
+
+def test_assign3():
+    assert py2vega(assign_func3, whitelist) == "'red'"
+
+
+def test_assign4():
+    assert py2vega(assign_func4, whitelist) == "'green'"
+
+
+def test_assign5():
+    assert py2vega(assign_func5, whitelist) == "['Hello', 'Hello']"
+
+
+def test_assign6():
+    assert py2vega(assign_func6, whitelist) == "'Hello'"