From 3f56620870a36b2c5fe1e402d7e8e0ad5cf7febf Mon Sep 17 00:00:00 2001
From: Elijah Ahianyo <elijahahianyo@gmail.com>
Date: Fri, 23 Jun 2023 19:56:05 +0000
Subject: [PATCH] Get Local Storage (#1234)

---
 pynecone/.templates/web/utils/state.js | 17 +++++++
 pynecone/__init__.py                   |  1 +
 pynecone/compiler/compiler.py          |  1 +
 pynecone/vars.py                       | 23 +++++++++
 tests/test_var.py                      | 69 +++++++++++++++++++++++++-
 5 files changed, 110 insertions(+), 1 deletion(-)

diff --git a/pynecone/.templates/web/utils/state.js b/pynecone/.templates/web/utils/state.js
index ec6ebe8f3..7657ef305 100644
--- a/pynecone/.templates/web/utils/state.js
+++ b/pynecone/.templates/web/utils/state.js
@@ -76,6 +76,23 @@ export const applyDelta = (state, delta) => {
   }
 };
 
+
+/**
+ * Get all local storage items in a key-value object.
+ * @returns object of items in local storage.
+ */
+export const getAllLocalStorageItems = () => {
+  var localStorageItems = {};
+
+  for (var i = 0, len = localStorage.length; i < len; i++) {
+    var key = localStorage.key(i);
+    localStorageItems[key] = localStorage.getItem(key);
+  }
+
+  return localStorageItems;
+}
+
+
 /**
  * Send an event to the server.
  * @param event The event to send.
diff --git a/pynecone/__init__.py b/pynecone/__init__.py
index e4ebf698a..d2d920807 100644
--- a/pynecone/__init__.py
+++ b/pynecone/__init__.py
@@ -38,3 +38,4 @@ from .style import color_mode as color_mode
 from .style import toggle_color_mode as toggle_color_mode
 from .vars import Var as Var
 from .vars import cached_var as cached_var
+from .vars import get_local_storage as get_local_storage
diff --git a/pynecone/compiler/compiler.py b/pynecone/compiler/compiler.py
index b456f5339..2ef27b50d 100644
--- a/pynecone/compiler/compiler.py
+++ b/pynecone/compiler/compiler.py
@@ -30,6 +30,7 @@ DEFAULT_IMPORTS: imports.ImportDict = {
         ImportVar(tag="preventDefault"),
         ImportVar(tag="refs"),
         ImportVar(tag="getRefValue"),
+        ImportVar(tag="getAllLocalStorageItems"),
     },
     "": {ImportVar(tag="focus-visible/dist/focus-visible")},
     "@chakra-ui/react": {
diff --git a/pynecone/vars.py b/pynecone/vars.py
index c40ab9b45..229169a1c 100644
--- a/pynecone/vars.py
+++ b/pynecone/vars.py
@@ -1133,3 +1133,26 @@ class ImportVar(Base):
             The hash of the var.
         """
         return hash((self.tag, self.is_default, self.alias))
+
+
+def get_local_storage(key: Optional[Union[Var, str]] = None) -> BaseVar:
+    """Provide a base var as payload to get local storage item(s).
+
+    Args:
+        key: Key to obtain value in the local storage.
+
+    Returns:
+        A BaseVar of the local storage method/function to call.
+
+    Raises:
+        TypeError:  if the wrong key type is provided.
+    """
+    if key:
+        if not (isinstance(key, Var) and key.type_ == str) and not isinstance(key, str):
+            type_ = type(key) if not isinstance(key, Var) else key.type_
+            raise TypeError(
+                f"Local storage keys can only be of type `str` or `var` of type `str`. Got `{type_}` instead."
+            )
+        key = key.full_name if isinstance(key, Var) else format.wrap(key, "'")
+        return BaseVar(name=f"localStorage.getItem({key})", type_=str)
+    return BaseVar(name="getAllLocalStorageItems()", type_=Dict)
diff --git a/tests/test_var.py b/tests/test_var.py
index a140a5168..87097d1a5 100644
--- a/tests/test_var.py
+++ b/tests/test_var.py
@@ -6,7 +6,15 @@ import pytest
 
 from pynecone.base import Base
 from pynecone.state import State
-from pynecone.vars import BaseVar, ComputedVar, ImportVar, PCDict, PCList, Var
+from pynecone.vars import (
+    BaseVar,
+    ComputedVar,
+    ImportVar,
+    PCDict,
+    PCList,
+    Var,
+    get_local_storage,
+)
 
 test_vars = [
     BaseVar(name="prop1", type_=int),
@@ -19,6 +27,12 @@ test_vars = [
 test_import_vars = [ImportVar(tag="DataGrid"), ImportVar(tag="DataGrid", alias="Grid")]
 
 
+class BaseState(State):
+    """A Test State."""
+
+    val: str = "key"
+
+
 @pytest.fixture
 def TestObj():
     class TestObj(Base):
@@ -394,3 +408,56 @@ def test_import_var(import_var, expected):
         expected: expected name
     """
     assert import_var.name == expected
+
+
+@pytest.mark.parametrize(
+    "key, expected",
+    [
+        ("test_key", BaseVar(name="localStorage.getItem('test_key')", type_=str)),
+        (
+            BaseVar(name="key_var", type_=str),
+            BaseVar(name="localStorage.getItem(key_var)", type_=str),
+        ),
+        (
+            BaseState.val,
+            BaseVar(name="localStorage.getItem(base_state.val)", type_=str),
+        ),
+        (None, BaseVar(name="getAllLocalStorageItems()", type_=Dict)),
+    ],
+)
+def test_get_local_storage(key, expected):
+    """Test that the right BaseVar is return when get_local_storage is called.
+
+    Args:
+        key: Local storage key.
+        expected: expected BaseVar.
+
+    """
+    local_storage = get_local_storage(key)
+    assert local_storage.name == expected.name
+    assert local_storage.type_ == expected.type_
+
+
+@pytest.mark.parametrize(
+    "key",
+    [
+        ["list", "values"],
+        {"name": "dict"},
+        10,
+        BaseVar(name="key_var", type_=List),
+        BaseVar(name="key_var", type_=Dict[str, str]),
+    ],
+)
+def test_get_local_storage_raise_error(key):
+    """Test that a type error is thrown when the wrong key type is provided.
+
+    Args:
+        key: Local storage key.
+    """
+    with pytest.raises(TypeError) as err:
+        get_local_storage(key)
+    type_ = type(key) if not isinstance(key, Var) else key.type_
+    assert (
+        err.value.args[0]
+        == f"Local storage keys can only be of type `str` or `var` of type `str`. Got `{type_}` instead."
+    )