Add custom node types support

examples/nodes.py

@@ -1,11 +1,12 @@
 from trame.app import TrameApp
 from trame.ui.vuetify3 import SinglePageLayout
+from trame.widgets.html import H4, P
 from trame.widgets.vuetify3 import (
     VBtn,
     VIcon,
     VSelect,
 )
-from trame_flow.widgets.flow import Node, NodeEditor, create_node
+from trame_flow.widgets.flow import CustomNode, Node, NodeEditor, create_node
 
 
 class Example(TrameApp):
@@ -26,6 +27,7 @@ def add_node(self):
                 y=(self.next_node_id // 10) * 100,
                 type=self.state.node_type,
                 label=f"Node {self.next_node_id}",
+                data={"subtitle": "subtitle"},
             )
         )
         self.next_node_id += 1
@@ -42,7 +44,7 @@ def build_ui(self):
             with layout.toolbar:
                 VSelect(
                     label="Node type",
-                    items=("['default', 'input', 'output', 'text']",),
+                    items=("['default', 'input', 'output', 'text', 'title']",),
                     v_model=("node_type", "default"),
                     density="compact",
                     hide_details="true",
@@ -59,7 +61,12 @@ def build_ui(self):
                 ):
                     VIcon("mdi-minus")
 
-            self.vueflow = NodeEditor()
+            with NodeEditor() as self.vueflow:
+                with CustomNode("title"):
+                    H4("{{props.data.label}}")
+                    P("{{props.data.subtitle}}")
+                with CustomNode("text", var_name="nodeProps"):
+                    P("{{nodeProps.data.label}}")
 
             def on_graph_change(nodes, edges):
                 with self.state:

src/trame_flow/widgets/flow.py

@@ -1,7 +1,7 @@
 from ast import literal_eval
 from typing import Callable, Literal, Optional, Union
 
-from trame_client.widgets.core import AbstractElement
+from trame_client.widgets.core import AbstractElement, Template
 from typing_extensions import NotRequired, TypedDict
 
 from .. import module
@@ -16,6 +16,7 @@ def __init__(self, _elem_name, children=None, **kwargs):
 
 __all__ = [
     "DEFAULT_EXTENT",
+    "CustomNode",
     "Dimensions",
     "Edge",
     "EdgeMarkerType",
@@ -49,7 +50,7 @@ class Dimensions(TypedDict):
 Extent = Union[Literal["parent"], list[list[float]]]
 DEFAULT_EXTENT = [[float("-inf"), float("-inf")], [float("+inf"), float("+inf")]]
 
-NodeType = Literal["default", "input", "output", "text"]
+NodeType = Literal["default", "input", "output"] | str
 
 
 Node = TypedDict(
@@ -114,8 +115,8 @@ def create_node(
         node["style"] = style
     if data:
         node["data"] = node["data"] | data
-    # set "text" node default css
-    if type == "text":
+    # set default node style for custom node
+    if type not in ["default", "input", "output"]:
         node["class"] = "vue-flow__node-default"
     return node
 
@@ -200,6 +201,19 @@ class Graph(TypedDict):
     edges: list[Edge]
 
 
+class CustomNode(Template):
+    def __init__(self, type: str, var_name: str = "props", **kwargs):
+        """Shortcut to define a custom node's HTML.
+        This is equivalent to `Template(raw_attrs=["v-slot:node-myType=myProps"])`
+
+        :param type: Name of the custom node type.
+        :type type: str
+        :param var_name: Name of the node properties variable. (Default = "props")
+        :type var_name:  str
+        """
+        super().__init__(raw_attrs=[f"v-slot:node-{type}={var_name}"], **kwargs)
+
+
 class NodeEditor(HtmlElement):
     """
     Node Editor based on VueFlow.

vue-components/src/components/NodeEditor.js

@@ -46,8 +46,9 @@ export default {
     <vue-flow elevate-edges-on-select="true" fit-view-on-init="true">
       <background :variant="backgroundPatternVariant" :pattern-color="backgroundPatternColor" :size="backgroundPatternSize" :gap="backgroundPatternGap" />
 
-      <template #node-text="nodeProps">
-        {{nodeProps.data.label}}
+      <!-- Forward all slots for custom nodes -->
+      <template v-for="(_, slotName) in $slots" v-slot:[slotName]="slotProps">
+        <slot :name="slotName" v-bind="slotProps ?? {}" />
       </template>
 
       <controls v-if="showControls" />