Implement var slicing (#289)

Author: picklelo

pynecone/compiler/utils.py

@@ -2,7 +2,7 @@
 
 import json
 import os
-from typing import Dict, List, Set, Tuple, Type
+from typing import Dict, List, Optional, Set, Tuple, Type
 
 from pynecone import constants, utils
 from pynecone.compiler import templates
@@ -306,13 +306,14 @@ def write_page(path: str, code: str):
         f.write(code)
 
 
-def empty_dir(path, keep_files=[]):
-    """Remove all files and folders in a directory except for the kept file- or foldernames.
+def empty_dir(path: str, keep_files: Optional[List[str]] = None):
+    """Remove all files and folders in a directory except for the keep_files.
 
     Args:
-        path (str): The path to the directory that will be emptied
-        keep_files (list, optional): List of filenames or foldernames that will not be deleted. Defaults to [].
+        path: The path to the directory that will be emptied
+        keep_files: List of filenames or foldernames that will not be deleted.
     """
+    keep_files = keep_files or []
     directory_contents = os.listdir(path)
     for element in directory_contents:
         if element not in keep_files:

pynecone/constants.py

@@ -2,7 +2,6 @@
 
 import os
 import re
-
 from enum import Enum
 from types import SimpleNamespace
 

pynecone/utils.py

@@ -15,7 +15,7 @@
 import sys
 from collections import defaultdict
 from pathlib import Path
-from subprocess import PIPE, DEVNULL, STDOUT
+from subprocess import DEVNULL, PIPE, STDOUT
 from types import ModuleType
 from typing import _GenericAlias  # type: ignore
 from typing import (

pynecone/var.py

@@ -138,33 +138,64 @@ def __getitem__(self, i: Any) -> Var:
         Raises:
             TypeError: If the var is not indexable.
         """
+        # Indexing is only supported for lists, dicts, and dataframes.
+        if not (
+            utils._issubclass(self.type_, Union[List, Dict])
+            or utils.is_dataframe(self.type_)
+        ):
+            raise TypeError(
+                f"Var {self.name} of type {self.type_} does not support indexing."
+            )
+
         # The type of the indexed var.
-        type_ = str
+        type_ = Any
 
         # Convert any vars to local vars.
         if isinstance(i, Var):
             i = BaseVar(name=i.name, type_=i.type_, state=i.state, is_local=True)
 
+        # Handle list indexing.
         if utils._issubclass(self.type_, List):
-            assert isinstance(
-                i, utils.get_args(Union[int, Var])
-            ), "Index must be an integer."
+            # List indices must be ints, slices, or vars.
+            if not isinstance(i, utils.get_args(Union[int, slice, Var])):
+                raise TypeError("Index must be an integer.")
+
+            # Handle slices first.
+            if isinstance(i, slice):
+                # Get the start and stop indices.
+                start = i.start or 0
+                stop = i.stop or "undefined"
+
+                # Use the slice function.
+                return BaseVar(
+                    name=f"{self.name}.slice({start}, {stop})",
+                    type_=self.type_,
+                    state=self.state,
+                )
+
+            # Get the type of the indexed var.
             if utils.is_generic_alias(self.type_):
                 type_ = utils.get_args(self.type_)[0]
             else:
                 type_ = Any
-        elif utils._issubclass(self.type_, Dict) or utils.is_dataframe(self.type_):
-            if isinstance(i, str):
-                i = utils.wrap(i, '"')
-            if utils.is_generic_alias(self.type_):
-                type_ = utils.get_args(self.type_)[1]
-            else:
-                type_ = Any
-        else:
-            raise TypeError(
-                f"Var {self.name} of type {self.type_} does not support indexing."
+
+            # Use `at` to support negative indices.
+            return BaseVar(
+                name=f"{self.name}.at({i})",
+                type_=type_,
+                state=self.state,
             )
 
+        # Dictionary / dataframe indexing.
+        # Get the type of the indexed var.
+        if isinstance(i, str):
+            i = utils.wrap(i, '"')
+        if utils.is_generic_alias(self.type_):
+            type_ = utils.get_args(self.type_)[1]
+        else:
+            type_ = Any
+
+        # Use normal indexing here.
         return BaseVar(
             name=f"{self.name}[{i}]",
             type_=type_,
@@ -621,6 +652,7 @@ class BaseVar(Var, Base):
     # Whether this is a local javascript variable.
     is_local: bool = False
 
+    # Whether this var is a raw string.
     is_string: bool = False
 
     def __hash__(self) -> int:

tests/test_state.py

@@ -251,10 +251,10 @@ def test_default_setters(test_state):
 def test_class_indexing_with_vars():
     """Test that we can index into a state var with another var."""
     prop = TestState.array[TestState.num1]
-    assert str(prop) == "{test_state.array[test_state.num1]}"
+    assert str(prop) == "{test_state.array.at(test_state.num1)}"
 
     prop = TestState.mapping["a"][TestState.num1]
-    assert str(prop) == '{test_state.mapping["a"][test_state.num1]}'
+    assert str(prop) == '{test_state.mapping["a"].at(test_state.num1)}'
 
 
 def test_class_attributes():

tests/test_var.py

@@ -1,3 +1,5 @@
+from typing import Dict, List
+
 import pytest
 
 from pynecone.base import Base
@@ -135,18 +137,18 @@ def test_create(value, expected):
         assert prop.equals(expected)  # type: ignore
 
 
+def v(value) -> Var:
+    val = Var.create(value)
+    assert val is not None
+    return val
+
+
 def test_basic_operations(TestObj):
     """Test the var operations.
 
     Args:
         TestObj: The test object.
     """
-
-    def v(value) -> Var:
-        val = Var.create(value)
-        assert val is not None
-        return val
-
     assert str(v(1) == v(2)) == "{(1 == 2)}"
     assert str(v(1) != v(2)) == "{(1 != 2)}"
     assert str(v(1) < v(2)) == "{(1 < 2)}"
@@ -162,8 +164,46 @@ def v(value) -> Var:
     assert str(v(1) ** v(2)) == "{Math.pow(1 , 2)}"
     assert str(v(1) & v(2)) == "{(1 && 2)}"
     assert str(v(1) | v(2)) == "{(1 || 2)}"
-    assert str(v([1, 2, 3])[v(0)]) == "{[1, 2, 3][0]}"
+    assert str(v([1, 2, 3])[v(0)]) == "{[1, 2, 3].at(0)}"
     assert str(v({"a": 1, "b": 2})["a"]) == '{{"a": 1, "b": 2}["a"]}'
     assert (
         str(BaseVar(name="foo", state="state", type_=TestObj).bar) == "{state.foo.bar}"
     )
+    assert str(abs(v(1))) == "{Math.abs(1)}"
+    assert str(v([1, 2, 3]).length()) == "{[1, 2, 3].length}"
+
+
+def test_var_indexing_lists():
+    """Test that we can index into list vars."""
+    lst = BaseVar(name="lst", type_=List[int])
+
+    # Test basic indexing.
+    assert str(lst[0]) == "{lst.at(0)}"
+    assert str(lst[1]) == "{lst.at(1)}"
+
+    # Test negative indexing.
+    assert str(lst[-1]) == "{lst.at(-1)}"
+
+    # Test non-integer indexing raises an error.
+    with pytest.raises(TypeError):
+        lst["a"]
+    with pytest.raises(TypeError):
+        lst[1.5]
+
+
+def test_var_list_slicing():
+    """Test that we can slice into list vars."""
+    lst = BaseVar(name="lst", type_=List[int])
+
+    assert str(lst[0:1]) == "{lst.slice(0, 1)}"
+    assert str(lst[:1]) == "{lst.slice(0, 1)}"
+    assert str(lst[0:]) == "{lst.slice(0, undefined)}"
+
+
+def test_dict_indexing():
+    """Test that we can index into dict vars."""
+    dct = BaseVar(name="dct", type_=Dict[str, int])
+
+    # Check correct indexing.
+    assert str(dct["a"]) == '{dct["a"]}'
+    assert str(dct["asdf"]) == '{dct["asdf"]}'