From 8eea6ac8229b7dc460af1e2a4f51214ffd32e415 Mon Sep 17 00:00:00 2001
From: Nikhil Rao <nikhil@pynecone.io>
Date: Mon, 3 Apr 2023 17:38:56 -0700
Subject: [PATCH] Add component memoization (#770)

---
 pynecone/__init__.py               |  2 +-
 tests/components/test_component.py | 24 +++++++++++++++++++++++-
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/pynecone/__init__.py b/pynecone/__init__.py
index f645e3cb9..2a09f52f9 100644
--- a/pynecone/__init__.py
+++ b/pynecone/__init__.py
@@ -7,7 +7,7 @@ from . import el
 from .app import App, UploadFile
 from .base import Base
 from .components import *
-from .components.component import custom_component as component
+from .components.component import custom_component as memo
 from .components.graphing.victory import data
 from .config import Config
 from .constants import Env, Transports
diff --git a/tests/components/test_component.py b/tests/components/test_component.py
index 099ae7aff..100589a37 100644
--- a/tests/components/test_component.py
+++ b/tests/components/test_component.py
@@ -2,7 +2,8 @@ from typing import Dict, List, Type
 
 import pytest
 
-from pynecone.components.component import Component, CustomComponent
+import pynecone as pc
+from pynecone.components.component import Component, CustomComponent, custom_component
 from pynecone.components.layout.box import Box
 from pynecone.event import EVENT_ARG, EVENT_TRIGGERS, EventHandler
 from pynecone.state import State
@@ -311,6 +312,27 @@ def test_custom_component_hash(my_component):
     assert {component1, component2} == {component1}
 
 
+def test_custom_component_wrapper():
+    """Test that the wrapper of a custom component is correct."""
+
+    @custom_component
+    def my_component(width: Var[int], color: Var[str]):
+        return pc.box(
+            width=width,
+            color=color,
+        )
+
+    ccomponent = my_component(
+        pc.text("child"), width=Var.create(1), color=Var.create("red")
+    )
+    assert isinstance(ccomponent, CustomComponent)
+    assert len(ccomponent.children) == 1
+    assert isinstance(ccomponent.children[0], pc.Text)
+
+    component = ccomponent.get_component()
+    assert isinstance(component, Box)
+
+
 def test_invalid_event_handler_args(component2, TestState):
     """Test that an invalid event handler raises an error.