From 18c715670ad8768b2fa0d307266e2da5827ecc70 Mon Sep 17 00:00:00 2001
From: Elijah Ahianyo <elijahahianyo@gmail.com>
Date: Fri, 5 May 2023 05:11:01 +0000
Subject: [PATCH] Revamp Imports (#926)

---
 pynecone/compiler/compiler.py                | 29 ++++++-----
 pynecone/compiler/utils.py                   | 13 ++---
 pynecone/components/base/head.py             |  2 +
 pynecone/components/component.py             | 37 ++++++++------
 pynecone/components/datadisplay/code.py      |  6 ++-
 pynecone/components/datadisplay/datatable.py | 18 +++----
 pynecone/components/forms/input.py           |  6 +--
 pynecone/components/forms/numberinput.py     |  2 +-
 pynecone/components/forms/select.py          |  2 +-
 pynecone/components/forms/slider.py          |  2 +-
 pynecone/components/forms/textarea.py        |  2 +-
 pynecone/components/forms/upload.py          |  2 +
 pynecone/components/navigation/nextlink.py   |  2 +
 pynecone/components/typography/markdown.py   | 30 ++++++-----
 pynecone/el/elements/__init__.py             |  2 +-
 pynecone/utils/imports.py                    |  4 +-
 pynecone/var.py                              | 31 ++++++++++-
 tests/compiler/test_compiler.py              | 54 ++++++++++++++++----
 tests/components/test_component.py           | 13 +++--
 tests/test_var.py                            | 24 ++++++++-
 20 files changed, 192 insertions(+), 89 deletions(-)

diff --git a/pynecone/compiler/compiler.py b/pynecone/compiler/compiler.py
index a35f4dab1..c54621279 100644
--- a/pynecone/compiler/compiler.py
+++ b/pynecone/compiler/compiler.py
@@ -11,20 +11,25 @@ from pynecone.components.component import Component, CustomComponent
 from pynecone.state import State
 from pynecone.style import Style
 from pynecone.utils import imports, path_ops
+from pynecone.var import ImportVar
 
 # Imports to be included in every Pynecone app.
 DEFAULT_IMPORTS: imports.ImportDict = {
-    "react": {"useEffect", "useRef", "useState"},
-    "next/router": {"useRouter"},
-    f"/{constants.STATE_PATH}": {
-        "connect",
-        "updateState",
-        "uploadFiles",
-        "E",
-        "isTrue",
+    "react": {
+        ImportVar(tag="useEffect"),
+        ImportVar(tag="useRef"),
+        ImportVar(tag="useState"),
     },
-    "": {"focus-visible/dist/focus-visible"},
-    "@chakra-ui/react": {constants.USE_COLOR_MODE},
+    "next/router": {ImportVar(tag="useRouter")},
+    f"/{constants.STATE_PATH}": {
+        ImportVar(tag="connect"),
+        ImportVar(tag="updateState"),
+        ImportVar(tag="uploadFiles"),
+        ImportVar(tag="E"),
+        ImportVar(tag="isTrue"),
+    },
+    "": {ImportVar(tag="focus-visible/dist/focus-visible")},
+    "@chakra-ui/react": {ImportVar(tag=constants.USE_COLOR_MODE)},
 }
 
 
@@ -91,8 +96,8 @@ def _compile_components(components: Set[CustomComponent]) -> str:
         The compiled components.
     """
     imports = {
-        "react": {"memo"},
-        f"/{constants.STATE_PATH}": {"E", "isTrue"},
+        "react": {ImportVar(tag="memo")},
+        f"/{constants.STATE_PATH}": {ImportVar(tag="E"), ImportVar(tag="isTrue")},
     }
     component_defs = []
 
diff --git a/pynecone/compiler/utils.py b/pynecone/compiler/utils.py
index d6b80f3a7..a7dba4f61 100644
--- a/pynecone/compiler/utils.py
+++ b/pynecone/compiler/utils.py
@@ -25,12 +25,13 @@ from pynecone.event import get_hydrate_event
 from pynecone.state import State
 from pynecone.style import Style
 from pynecone.utils import format, imports, path_ops
+from pynecone.var import ImportVar
 
 # To re-export this function.
 merge_imports = imports.merge_imports
 
 
-def compile_import_statement(lib: str, fields: Set[str]) -> str:
+def compile_import_statement(lib: str, fields: Set[ImportVar]) -> str:
     """Compile an import statement.
 
     Args:
@@ -41,16 +42,12 @@ def compile_import_statement(lib: str, fields: Set[str]) -> str:
         The compiled import statement.
     """
     # Check for default imports.
-    defaults = {
-        field
-        for field in fields
-        if field.lower() == lib.lower().replace("-", "").replace("/", "")
-    }
+    defaults = {field for field in fields if field.is_default}
     assert len(defaults) < 2
 
     # Get the default import, and the specific imports.
-    default = next(iter(defaults), "")
-    rest = fields - defaults
+    default = next(iter({field.name for field in defaults}), "")
+    rest = {field.name for field in fields - defaults}
     return templates.format_import(lib=lib, default=default, rest=rest)
 
 
diff --git a/pynecone/components/base/head.py b/pynecone/components/base/head.py
index 314f78070..dcaaff166 100644
--- a/pynecone/components/base/head.py
+++ b/pynecone/components/base/head.py
@@ -13,3 +13,5 @@ class Head(NextHeadLib):
     """Head Component."""
 
     tag = "NextHead"
+
+    is_default = True
diff --git a/pynecone/components/component.py b/pynecone/components/component.py
index b5c56b350..758211913 100644
--- a/pynecone/components/component.py
+++ b/pynecone/components/component.py
@@ -22,7 +22,7 @@ from pynecone.event import (
 )
 from pynecone.style import Style
 from pynecone.utils import format, imports, path_ops, types
-from pynecone.var import BaseVar, Var
+from pynecone.var import BaseVar, ImportVar, Var
 
 
 class Component(Base, ABC):
@@ -43,6 +43,12 @@ class Component(Base, ABC):
     # The tag to use when rendering the component.
     tag: Optional[str] = None
 
+    # The alias for the tag.
+    alias: Optional[str] = None
+
+    # Whether the import is default or named.
+    is_default: Optional[bool] = False
+
     # A unique key for the component.
     key: Any = None
 
@@ -275,15 +281,6 @@ class Component(Base, ABC):
         """
         return {}
 
-    @classmethod
-    def get_alias(cls) -> Optional[str]:
-        """Get the alias for the component.
-
-        Returns:
-            The alias.
-        """
-        return None
-
     def __repr__(self) -> str:
         """Represent the component in React.
 
@@ -307,9 +304,10 @@ class Component(Base, ABC):
             The tag to render.
         """
         # Create the base tag.
-        alias = self.get_alias()
-        name = alias if alias is not None else self.tag
-        tag = Tag(name=name, special_props=self.special_props)
+        tag = Tag(
+            name=self.tag if not self.alias else self.alias,
+            special_props=self.special_props,
+        )
 
         # Add component props to the tag.
         props = {attr: getattr(self, attr) for attr in self.get_props()}
@@ -445,9 +443,7 @@ class Component(Base, ABC):
 
     def _get_imports(self) -> imports.ImportDict:
         if self.library is not None and self.tag is not None:
-            alias = self.get_alias()
-            tag = self.tag if alias is None else " as ".join([self.tag, alias])
-            return {self.library: {tag}}
+            return {self.library: {self.import_var}}
         return {}
 
     def get_imports(self) -> imports.ImportDict:
@@ -521,6 +517,15 @@ class Component(Base, ABC):
             custom_components |= child.get_custom_components(seen=seen)
         return custom_components
 
+    @property
+    def import_var(self):
+        """The tag to import.
+
+        Returns:
+            An import var.
+        """
+        return ImportVar(tag=self.tag, is_default=self.is_default, alias=self.alias)
+
     def is_full_control(self, kwargs: dict) -> bool:
         """Return if the component is fully controlled input.
 
diff --git a/pynecone/components/datadisplay/code.py b/pynecone/components/datadisplay/code.py
index 7ef91759f..d7256f48a 100644
--- a/pynecone/components/datadisplay/code.py
+++ b/pynecone/components/datadisplay/code.py
@@ -6,7 +6,7 @@ from pynecone.components.component import Component
 from pynecone.components.libs.chakra import ChakraComponent
 from pynecone.style import Style
 from pynecone.utils import imports
-from pynecone.var import Var
+from pynecone.var import ImportVar, Var
 
 # Path to the prism styles.
 PRISM_STYLES_PATH = "/styles/code/prism"
@@ -19,6 +19,8 @@ class CodeBlock(Component):
 
     tag = "Prism"
 
+    is_default = True
+
     # The theme to use ("light" or "dark").
     theme: Var[str]
 
@@ -44,7 +46,7 @@ class CodeBlock(Component):
         merged_imports = super()._get_imports()
         if self.theme is not None:
             merged_imports = imports.merge_imports(
-                merged_imports, {PRISM_STYLES_PATH: {self.theme.name}}
+                merged_imports, {PRISM_STYLES_PATH: {ImportVar(tag=self.theme.name)}}
             )
         return merged_imports
 
diff --git a/pynecone/components/datadisplay/datatable.py b/pynecone/components/datadisplay/datatable.py
index 2ec6baa91..3f3277b42 100644
--- a/pynecone/components/datadisplay/datatable.py
+++ b/pynecone/components/datadisplay/datatable.py
@@ -1,11 +1,11 @@
 """Table components."""
 
-from typing import Any, List, Optional
+from typing import Any, List
 
 from pynecone.components.component import Component
 from pynecone.components.tags import Tag
 from pynecone.utils import format, imports, types
-from pynecone.var import BaseVar, ComputedVar, Var
+from pynecone.var import BaseVar, ComputedVar, ImportVar, Var
 
 
 class Gridjs(Component):
@@ -19,6 +19,8 @@ class DataTable(Gridjs):
 
     tag = "Grid"
 
+    alias = "DataTableGrid"
+
     # The data to display. Either a list of lists or a pandas dataframe.
     data: Any
 
@@ -38,15 +40,6 @@ class DataTable(Gridjs):
     # Enable pagination.
     pagination: Var[bool]
 
-    @classmethod
-    def get_alias(cls) -> Optional[str]:
-        """Get the alias for the component.
-
-        Returns:
-            The alias.
-        """
-        return "DataTableGrid"
-
     @classmethod
     def create(cls, *children, **props):
         """Create a datatable component.
@@ -102,7 +95,8 @@ class DataTable(Gridjs):
 
     def _get_imports(self) -> imports.ImportDict:
         return imports.merge_imports(
-            super()._get_imports(), {"": {"gridjs/dist/theme/mermaid.css"}}
+            super()._get_imports(),
+            {"": {ImportVar(tag="gridjs/dist/theme/mermaid.css")}},
         )
 
     def _render(self) -> Tag:
diff --git a/pynecone/components/forms/input.py b/pynecone/components/forms/input.py
index dc07e9b1d..1561152f9 100644
--- a/pynecone/components/forms/input.py
+++ b/pynecone/components/forms/input.py
@@ -5,7 +5,7 @@ from typing import Dict
 from pynecone.components.component import EVENT_ARG
 from pynecone.components.libs.chakra import ChakraComponent
 from pynecone.utils import imports
-from pynecone.var import Var
+from pynecone.var import ImportVar, Var
 
 
 class Input(ChakraComponent):
@@ -13,7 +13,7 @@ class Input(ChakraComponent):
 
     tag = "Input"
 
-    # State var to bind the the input.
+    # State var to bind the input.
     value: Var[str]
 
     # The default value of the input.
@@ -52,7 +52,7 @@ class Input(ChakraComponent):
     def _get_imports(self) -> imports.ImportDict:
         return imports.merge_imports(
             super()._get_imports(),
-            {"/utils/state": {"set_val"}},
+            {"/utils/state": {ImportVar(tag="set_val")}},
         )
 
     @classmethod
diff --git a/pynecone/components/forms/numberinput.py b/pynecone/components/forms/numberinput.py
index c237d4fed..91e85865f 100644
--- a/pynecone/components/forms/numberinput.py
+++ b/pynecone/components/forms/numberinput.py
@@ -13,7 +13,7 @@ class NumberInput(ChakraComponent):
 
     tag = "NumberInput"
 
-    # State var to bind the the input.
+    # State var to bind the input.
     value: Var[int]
 
     # If true, the input's value will change based on mouse wheel.
diff --git a/pynecone/components/forms/select.py b/pynecone/components/forms/select.py
index 691dd882b..0409fa5b1 100644
--- a/pynecone/components/forms/select.py
+++ b/pynecone/components/forms/select.py
@@ -15,7 +15,7 @@ class Select(ChakraComponent):
 
     tag = "Select"
 
-    # State var to bind the the select.
+    # State var to bind the select.
     value: Var[str]
 
     # The default value of the select.
diff --git a/pynecone/components/forms/slider.py b/pynecone/components/forms/slider.py
index 4f9af579e..5745678c8 100644
--- a/pynecone/components/forms/slider.py
+++ b/pynecone/components/forms/slider.py
@@ -13,7 +13,7 @@ class Slider(ChakraComponent):
 
     tag = "Slider"
 
-    # State var to bind the the input.
+    # State var to bind the input.
     value: Var[int]
 
     # The color scheme.
diff --git a/pynecone/components/forms/textarea.py b/pynecone/components/forms/textarea.py
index ae9ee38bb..b5f41aca0 100644
--- a/pynecone/components/forms/textarea.py
+++ b/pynecone/components/forms/textarea.py
@@ -12,7 +12,7 @@ class TextArea(ChakraComponent):
 
     tag = "Textarea"
 
-    # State var to bind the the input.
+    # State var to bind the input.
     value: Var[str]
 
     # The default value of the textarea.
diff --git a/pynecone/components/forms/upload.py b/pynecone/components/forms/upload.py
index 1d4f65f74..2ad327001 100644
--- a/pynecone/components/forms/upload.py
+++ b/pynecone/components/forms/upload.py
@@ -18,6 +18,8 @@ class Upload(Component):
 
     tag = "ReactDropzone"
 
+    is_default = True
+
     # The list of accepted file types. This should be a dictionary of MIME types as keys and array of file formats as
     # values.
     # supported MIME types: https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/MIME_types/Common_types
diff --git a/pynecone/components/navigation/nextlink.py b/pynecone/components/navigation/nextlink.py
index a4c78f1de..08970836f 100644
--- a/pynecone/components/navigation/nextlink.py
+++ b/pynecone/components/navigation/nextlink.py
@@ -11,6 +11,8 @@ class NextLink(Component):
 
     tag = "NextLink"
 
+    is_default = True
+
     # The page to link to.
     href: Var[str]
 
diff --git a/pynecone/components/typography/markdown.py b/pynecone/components/typography/markdown.py
index 8d50912d1..d526b5468 100644
--- a/pynecone/components/typography/markdown.py
+++ b/pynecone/components/typography/markdown.py
@@ -5,7 +5,7 @@ from typing import List, Union
 
 from pynecone.components.component import Component
 from pynecone.utils import types
-from pynecone.var import BaseVar, Var
+from pynecone.var import BaseVar, ImportVar, Var
 
 
 class Markdown(Component):
@@ -15,6 +15,8 @@ class Markdown(Component):
 
     tag = "ReactMarkdown"
 
+    is_default = True
+
     @classmethod
     def create(cls, *children, **props) -> Component:
         """Create a markdown component.
@@ -39,20 +41,20 @@ class Markdown(Component):
     def _get_imports(self):
         imports = super()._get_imports()
         imports["@chakra-ui/react"] = {
-            "Heading",
-            "Code",
-            "Text",
-            "Link",
-            "UnorderedList",
-            "OrderedList",
-            "ListItem",
+            ImportVar(tag="Heading"),
+            ImportVar(tag="Code"),
+            ImportVar(tag="Text"),
+            ImportVar(tag="Link"),
+            ImportVar(tag="UnorderedList"),
+            ImportVar(tag="OrderedList"),
+            ImportVar(tag="ListItem"),
         }
-        imports["react-syntax-highlighter"] = {"Prism"}
-        imports["remark-math"] = {"remarkMath"}
-        imports["remark-gfm"] = {"remarkGfm"}
-        imports["rehype-katex"] = {"rehypeKatex"}
-        imports["rehype-raw"] = {"rehypeRaw"}
-        imports[""] = {"katex/dist/katex.min.css"}
+        imports["react-syntax-highlighter"] = {ImportVar(tag="Prism", is_default=True)}
+        imports["remark-math"] = {ImportVar(tag="remarkMath", is_default=True)}
+        imports["remark-gfm"] = {ImportVar(tag="remarkGfm", is_default=True)}
+        imports["rehype-katex"] = {ImportVar(tag="rehypeKatex", is_default=True)}
+        imports["rehype-raw"] = {ImportVar(tag="rehypeRaw", is_default=True)}
+        imports[""] = {ImportVar(tag="katex/dist/katex.min.css")}
         return imports
 
     def _render(self):
diff --git a/pynecone/el/elements/__init__.py b/pynecone/el/elements/__init__.py
index 93acaf3cb..055dbebff 100644
--- a/pynecone/el/elements/__init__.py
+++ b/pynecone/el/elements/__init__.py
@@ -6,7 +6,7 @@ from pynecone.var import Var as PCVar
 
 
 class A(Element):  # noqa: E742
-    """Display the a element."""
+    """Display the 'a' element."""
 
     tag = "a"
 
diff --git a/pynecone/utils/imports.py b/pynecone/utils/imports.py
index e955e56cf..a745dac1a 100644
--- a/pynecone/utils/imports.py
+++ b/pynecone/utils/imports.py
@@ -3,7 +3,9 @@
 from collections import defaultdict
 from typing import Dict, Set
 
-ImportDict = Dict[str, Set[str]]
+from pynecone.var import ImportVar
+
+ImportDict = Dict[str, Set[ImportVar]]
 
 
 def merge_imports(*imports) -> ImportDict:
diff --git a/pynecone/var.py b/pynecone/var.py
index 74ae1c80c..c6b487437 100644
--- a/pynecone/var.py
+++ b/pynecone/var.py
@@ -34,7 +34,6 @@ from pynecone.utils import format, types
 if TYPE_CHECKING:
     from pynecone.state import State
 
-
 # Set of unique variable names.
 USED_VARIABLES = set()
 
@@ -1071,3 +1070,33 @@ class PCDict(dict):
         """
         super().__delitem__(*args, **kwargs)
         self._reassign_field()
+
+
+class ImportVar(Base):
+    """An import var."""
+
+    # The name of the import tag.
+    tag: Optional[str]
+
+    # whether the import is default or named.
+    is_default: Optional[bool] = False
+
+    # The tag alias.
+    alias: Optional[str] = None
+
+    @property
+    def name(self) -> str:
+        """The name of the import.
+
+        Returns:
+            The name(tag name with alias) of tag.
+        """
+        return self.tag if not self.alias else " as ".join([self.tag, self.alias])  # type: ignore
+
+    def __hash__(self) -> int:
+        """Define a hash function for the import var.
+
+        Returns:
+            The hash of the var.
+        """
+        return hash((self.tag, self.is_default, self.alias))
diff --git a/tests/compiler/test_compiler.py b/tests/compiler/test_compiler.py
index 31c2dae49..a1bab5eae 100644
--- a/tests/compiler/test_compiler.py
+++ b/tests/compiler/test_compiler.py
@@ -4,17 +4,34 @@ import pytest
 
 from pynecone.compiler import utils
 from pynecone.utils import imports
+from pynecone.var import ImportVar
 
 
 @pytest.mark.parametrize(
     "lib,fields,output",
     [
-        ("axios", {"axios"}, 'import axios from "axios"'),
-        ("axios", {"foo", "bar"}, 'import {bar, foo} from "axios"'),
-        ("axios", {"axios", "foo", "bar"}, 'import axios, {bar, foo} from "axios"'),
+        (
+            "axios",
+            {ImportVar(tag="axios", is_default=True)},
+            'import axios from "axios"',
+        ),
+        (
+            "axios",
+            {ImportVar(tag="foo"), ImportVar(tag="bar")},
+            'import {bar, foo} from "axios"',
+        ),
+        (
+            "axios",
+            {
+                ImportVar(tag="axios", is_default=True),
+                ImportVar(tag="foo"),
+                ImportVar(tag="bar"),
+            },
+            "import " "axios, " "{bar, " "foo} from " '"axios"',
+        ),
     ],
 )
-def test_compile_import_statement(lib: str, fields: Set[str], output: str):
+def test_compile_import_statement(lib: str, fields: Set[ImportVar], output: str):
     """Test the compile_import_statement function.
 
     Args:
@@ -29,15 +46,34 @@ def test_compile_import_statement(lib: str, fields: Set[str], output: str):
     "import_dict,output",
     [
         ({}, ""),
-        ({"axios": {"axios"}}, 'import axios from "axios"'),
-        ({"axios": {"foo", "bar"}}, 'import {bar, foo} from "axios"'),
         (
-            {"axios": {"axios", "foo", "bar"}, "react": {"react"}},
+            {"axios": {ImportVar(tag="axios", is_default=True)}},
+            'import axios from "axios"',
+        ),
+        (
+            {"axios": {ImportVar(tag="foo"), ImportVar(tag="bar")}},
+            'import {bar, foo} from "axios"',
+        ),
+        (
+            {
+                "axios": {
+                    ImportVar(tag="axios", is_default=True),
+                    ImportVar(tag="foo"),
+                    ImportVar(tag="bar"),
+                },
+                "react": {ImportVar(tag="react", is_default=True)},
+            },
             'import axios, {bar, foo} from "axios"\nimport react from "react"',
         ),
-        ({"": {"lib1.js", "lib2.js"}}, 'import "lib1.js"\nimport "lib2.js"'),
         (
-            {"": {"lib1.js", "lib2.js"}, "axios": {"axios"}},
+            {"": {ImportVar(tag="lib1.js"), ImportVar(tag="lib2.js")}},
+            'import "lib1.js"\nimport "lib2.js"',
+        ),
+        (
+            {
+                "": {ImportVar(tag="lib1.js"), ImportVar(tag="lib2.js")},
+                "axios": {ImportVar(tag="axios", is_default=True)},
+            },
             'import "lib1.js"\nimport "lib2.js"\nimport axios from "axios"',
         ),
     ],
diff --git a/tests/components/test_component.py b/tests/components/test_component.py
index 9ef5a0ac4..a3e0c4a3c 100644
--- a/tests/components/test_component.py
+++ b/tests/components/test_component.py
@@ -9,7 +9,7 @@ from pynecone.event import EVENT_ARG, EVENT_TRIGGERS, EventHandler
 from pynecone.state import State
 from pynecone.style import Style
 from pynecone.utils import imports
-from pynecone.var import Var
+from pynecone.var import ImportVar, Var
 
 
 @pytest.fixture
@@ -42,7 +42,7 @@ def component1() -> Type[Component]:
         number: Var[int]
 
         def _get_imports(self) -> imports.ImportDict:
-            return {"react": {"Component"}}
+            return {"react": {ImportVar(tag="Component")}}
 
         def _get_custom_code(self) -> str:
             return "console.log('component1')"
@@ -75,7 +75,7 @@ def component2() -> Type[Component]:
             }
 
         def _get_imports(self) -> imports.ImportDict:
-            return {"react-redux": {"connect"}}
+            return {"react-redux": {ImportVar(tag="connect")}}
 
         def _get_custom_code(self) -> str:
             return "console.log('component2')"
@@ -206,8 +206,11 @@ def test_get_imports(component1, component2):
     """
     c1 = component1.create()
     c2 = component2.create(c1)
-    assert c1.get_imports() == {"react": {"Component"}}
-    assert c2.get_imports() == {"react-redux": {"connect"}, "react": {"Component"}}
+    assert c1.get_imports() == {"react": {ImportVar(tag="Component")}}
+    assert c2.get_imports() == {
+        "react-redux": {ImportVar(tag="connect")},
+        "react": {ImportVar(tag="Component")},
+    }
 
 
 def test_get_custom_code(component1, component2):
diff --git a/tests/test_var.py b/tests/test_var.py
index c2aba06e5..4d2609dfc 100644
--- a/tests/test_var.py
+++ b/tests/test_var.py
@@ -4,7 +4,7 @@ import cloudpickle
 import pytest
 
 from pynecone.base import Base
-from pynecone.var import BaseVar, PCDict, PCList, Var
+from pynecone.var import BaseVar, ImportVar, PCDict, PCList, Var
 
 test_vars = [
     BaseVar(name="prop1", type_=int),
@@ -14,6 +14,8 @@ test_vars = [
     BaseVar(name="local2", type_=str, is_local=True),
 ]
 
+test_import_vars = [ImportVar(tag="DataGrid"), ImportVar(tag="DataGrid", alias="Grid")]
+
 
 @pytest.fixture
 def TestObj():
@@ -245,3 +247,23 @@ def test_pickleable_pc_dict():
 
     pickled_dict = cloudpickle.dumps(pc_dict)
     assert cloudpickle.loads(pickled_dict) == pc_dict
+
+
+@pytest.mark.parametrize(
+    "import_var,expected",
+    zip(
+        test_import_vars,
+        [
+            "DataGrid",
+            "DataGrid as Grid",
+        ],
+    ),
+)
+def test_import_var(import_var, expected):
+    """Test that the import var name is computed correctly.
+
+    Args:
+        import_var: The import var.
+        expected: expected name
+    """
+    assert import_var.name == expected