From 86a1e8a07e73dc468d5bcac0c2c96d96645c513d Mon Sep 17 00:00:00 2001
From: Masen Furer <m_github@0x26.net>
Date: Thu, 14 Mar 2024 16:24:12 -0700
Subject: [PATCH] Account for imports of custom components for frontend
 installation

Create new ImportVar that only install the package to avoid rendering the
imports in the page that contains the custom component itself.
---
 reflex/components/component.py | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/reflex/components/component.py b/reflex/components/component.py
index a81e94178..e33f4e26c 100644
--- a/reflex/components/component.py
+++ b/reflex/components/component.py
@@ -1265,6 +1265,9 @@ class CustomComponent(Component):
     # The props of the component.
     props: Dict[str, Any] = {}
 
+    # Props that reference other components.
+    component_props: Dict[str, Component] = {}
+
     def __init__(self, *args, **kwargs):
         """Initialize the custom component.
 
@@ -1302,6 +1305,7 @@ class CustomComponent(Component):
 
                 # Track hooks and imports associated with Component instances.
                 if base_value is not None and isinstance(value, Component):
+                    self.component_props[key] = value
                     value = base_value._replace(
                         merge_var_data=VarData(  # type: ignore
                             imports=value.get_imports(),
@@ -1335,6 +1339,25 @@ class CustomComponent(Component):
         """
         return hash(self.tag)
 
+    def _get_imports(self) -> Dict[str, List[ImportVar]]:
+        """Get the imports for the component.
+
+        This is needed because otherwise the imports for the component are not
+        installed during compile time, but they are rendered into the page.
+
+        Returns:
+            The imports for the component and any custom component props.
+        """
+        return imports.merge_imports(
+            super()._get_imports(),
+            # Sweep up any imports from CustomComponent props for frontend installation.
+            {
+                library: [ImportVar(tag=None, render=False, install=True)]
+                for comp in self.get_custom_components()
+                for library in comp.get_component(comp).get_imports()
+            },
+        )
+
     @classmethod
     def get_props(cls) -> Set[str]:
         """Get the props for the component.
@@ -1368,6 +1391,16 @@ class CustomComponent(Component):
             custom_components |= self.get_component(self).get_custom_components(
                 seen=seen
             )
+
+        # Fetch custom components from props as well.
+        for child_component in self.component_props.values():
+            if child_component.tag is None:
+                continue
+            if child_component.tag not in seen:
+                seen.add(child_component.tag)
+                if isinstance(child_component, CustomComponent):
+                    custom_components |= {child_component}
+                custom_components |= child_component.get_custom_components(seen=seen)
         return custom_components
 
     def _render(self) -> Tag: