From 8be411b81bada250d15fedc459bfd6fea09fb234 Mon Sep 17 00:00:00 2001 From: Nikhil Rao Date: Wed, 18 Jan 2023 17:53:04 -0800 Subject: [PATCH] Implement var slicing (#289) --- pynecone/compiler/utils.py | 11 +++---- pynecone/constants.py | 1 - pynecone/utils.py | 2 +- pynecone/var.py | 60 +++++++++++++++++++++++++++++--------- tests/test_state.py | 4 +-- tests/test_var.py | 54 +++++++++++++++++++++++++++++----- 6 files changed, 102 insertions(+), 30 deletions(-) diff --git a/pynecone/compiler/utils.py b/pynecone/compiler/utils.py index 17f2846c4..73753e437 100644 --- a/pynecone/compiler/utils.py +++ b/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: diff --git a/pynecone/constants.py b/pynecone/constants.py index aef215efb..a1a09dc6d 100644 --- a/pynecone/constants.py +++ b/pynecone/constants.py @@ -2,7 +2,6 @@ import os import re - from enum import Enum from types import SimpleNamespace diff --git a/pynecone/utils.py b/pynecone/utils.py index e48b1b2f1..39ad279eb 100644 --- a/pynecone/utils.py +++ b/pynecone/utils.py @@ -15,7 +15,7 @@ import subprocess 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 ( diff --git a/pynecone/var.py b/pynecone/var.py index 00ceca54c..d3b8e88d7 100644 --- a/pynecone/var.py +++ b/pynecone/var.py @@ -138,33 +138,64 @@ class Var(ABC): 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: diff --git a/tests/test_state.py b/tests/test_state.py index ec7394a99..0d076a3c8 100644 --- a/tests/test_state.py +++ b/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(): diff --git a/tests/test_var.py b/tests/test_var.py index 49eba20e6..74025b874 100644 --- a/tests/test_var.py +++ b/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 test_basic_operations(TestObj): 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"]}'