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')

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"`
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'"

pip install -e .[dev]
parse_log_pytest
requirements.txt
pytest
3.7
pytest --no-header -rA --tb=line --color=no -p no:cacheprovider -W ignore::DeprecationWarning
certifi @ file:///croot/certifi_1671487769961/work/certifi
exceptiongroup==1.2.2
importlib-metadata==6.7.0
iniconfig==2.0.0
packaging==24.0
pluggy==1.2.0
-e git+https://github.com/QuantStack/py2vega.git@7d7ba87e17eee93e3dffdcf16c7528320e267832#egg=py2vega
pytest==7.4.4
tomli==2.0.1
typing_extensions==4.7.1
zipp==3.15.0

name: py2vega
channels:
  - defaults
  - https://repo.anaconda.com/pkgs/main
  - https://repo.anaconda.com/pkgs/r
  - conda-forge
dependencies:
  - _libgcc_mutex=0.1=main
  - _openmp_mutex=5.1=1_gnu
  - ca-certificates=2025.2.25=h06a4308_0
  - certifi=2022.12.7=py37h06a4308_0
  - ld_impl_linux-64=2.40=h12ee557_0
  - libffi=3.4.4=h6a678d5_1
  - libgcc-ng=11.2.0=h1234567_1
  - libgomp=11.2.0=h1234567_1
  - libstdcxx-ng=11.2.0=h1234567_1
  - ncurses=6.4=h6a678d5_0
  - openssl=1.1.1w=h7f8727e_0
  - pip=22.3.1=py37h06a4308_0
  - python=3.7.16=h7a1cb2a_0
  - readline=8.2=h5eee18b_0
  - setuptools=65.6.3=py37h06a4308_0
  - sqlite=3.45.3=h5eee18b_0
  - tk=8.6.14=h39e8969_0
  - wheel=0.38.4=py37h06a4308_0
  - xz=5.6.4=h5eee18b_1
  - zlib=1.2.13=h5eee18b_1
  - pip:
      - exceptiongroup==1.2.2
      - importlib-metadata==6.7.0
      - iniconfig==2.0.0
      - packaging==24.0
      - pluggy==1.2.0
      - pytest==7.4.4
      - tomli==2.0.1
      - typing-extensions==4.7.1
      - zipp==3.15.0
prefix: /opt/conda/envs/py2vega

test/test_py2vega.py::test_assign1 test/test_py2vega.py::test_assign2 test/test_py2vega.py::test_assign3 test/test_py2vega.py::test_assign4 test/test_py2vega.py::test_assign5 test/test_py2vega.py::test_assign6
test/test_py2vega.py::test_nameconstant test/test_py2vega.py::test_num test/test_py2vega.py::test_str test/test_py2vega.py::test_tuple test/test_py2vega.py::test_list test/test_py2vega.py::test_dict test/test_py2vega.py::test_unary test/test_py2vega.py::test_binary test/test_py2vega.py::test_ternary test/test_py2vega.py::test_compare test/test_py2vega.py::test_function test/test_py2vega.py::test_whitelist test/test_py2vega.py::test_math test/test_py2vega.py::test_invalid1 test/test_py2vega.py::test_invalid2 test/test_py2vega.py::test_if_stmt
README.md py2vega/main.py
README.md py2vega/main.py