Repository: QuantStack/py2vega

Issue:
Support Python Slice node
We could map the `Slice` node to a Vega `slice` function call

Discussion:

Patch:
diff --git a/py2vega/main.py b/py2vega/main.py
index f153309..9ccecf1 100644
--- a/py2vega/main.py
+++ b/py2vega/main.py
@@ -240,6 +240,30 @@ class VegaExpressionVisitor(ast.NodeVisitor):
 
         raise NameError('name \'{}\' is not defined, only a subset of Python is supported'.format(func_name))
 
+    def visit_Subscript(self, node):
+        """Turn a Python Subscript node into a Vega-expression."""
+        value = self.visit(node.value)
+
+        if isinstance(node.slice, ast.Index):
+            return '{value}[{index}]'.format(
+                value=value,
+                index=self.visit(node.slice.value)
+            )
+
+        if isinstance(node.slice, ast.Slice):
+            if node.slice.step is not None:
+                raise Py2VegaSyntaxError('Unsupported step for {} node, only a subset of Python is supported'.format(
+                    node.slice.__class__.__name__))
+
+            args = [value, '0' if node.slice.lower is None else self.visit(node.slice.lower)]
+            if node.slice.upper is not None:
+                args += self.visit(node.slice.upper)
+
+            return 'slice({args})'.format(args=', '.join(args))
+
+        raise Py2VegaSyntaxError('Unsupported {} node, only a subset of Python is supported'.format(
+            node.slice.__class__.__name__))
+
     def visit_Attribute(self, node):
         """Turn a Python attribute expression into a Vega-expression."""
         return node.attr


Test patch:
diff --git a/test/test_py2vega.py b/test/test_py2vega.py
index 4fbbb57..cfaf5b0 100644
--- a/test/test_py2vega.py
+++ b/test/test_py2vega.py
@@ -122,6 +122,31 @@ def test_call():
     assert py2vega(code, whitelist) == 'length(value)'
 
 
+def test_subscript():
+    code = 'value[0]'
+    assert py2vega(code, whitelist) == 'value[0]'
+
+    code = '[34, 32][0 if value < 2 else 1]'
+    assert py2vega(code, whitelist) == '[34, 32][((value < 2) ? 0 : 1)]'
+
+    code = 'value[:2]'
+    assert py2vega(code, whitelist) == 'slice(value, 0, 2)'
+
+    code = 'value[1:]'
+    assert py2vega(code, whitelist) == 'slice(value, 1)'
+
+    code = 'value[:]'
+    assert py2vega(code, whitelist) == 'slice(value, 0)'
+
+    code = 'value[1:2]'
+    assert py2vega(code, whitelist) == 'slice(value, 1, 2)'
+
+    # Unsupported step parameter
+    code = 'value[::2]'
+    with pytest.raises(Py2VegaSyntaxError):
+        py2vega(code, whitelist)
+
+
 def foo(value):
     return 'red' if value < 150 else 'green'