support visualize node in the notebook (#43)

* widget support node
* update doc

aiida_worktree/__init__.py

@@ -3,6 +3,6 @@
 from .decorator import node, build_node
 
 
-__version__ = "0.1.5"
+__version__ = "0.1.6"
 
 __all__ = ["WorkTree", "Node", "node", "build_node"]

aiida_worktree/node.py

@@ -1,6 +1,7 @@
 from node_graph.node import Node as GraphNode
 from aiida_worktree.properties import property_pool
 from aiida_worktree.sockets import socket_pool
+from aiida_worktree.widget import NodeGraphWidget
 
 
 class Node(GraphNode):
@@ -22,6 +23,9 @@ def __init__(self, **kwargs):
         self.wait = None
         self.process = None
         self.pk = None
+        self._widget = NodeGraphWidget(
+            settings={"minmap": False}, style={"width": "40%", "height": "600px"}
+        )
 
     def to_dict(self):
         ndata = super().to_dict()
@@ -63,3 +67,11 @@ def from_dict(cls, data, node_pool=None):
     def reset(self):
         self.process = None
         self.state = "CREATED"
+
+    def _repr_mimebundle_(self, *args, **kwargs):
+        # if ipywdigets > 8.0.0, use _repr_mimebundle_ instead of _ipython_display_
+        self._widget.from_node(self)
+        if hasattr(self._widget, "_repr_mimebundle_"):
+            return self._widget._repr_mimebundle_(*args, **kwargs)
+        else:
+            return self._widget._ipython_display_(*args, **kwargs)

aiida_worktree/widget/js/default_rete.ts

@@ -47,7 +47,7 @@ interface NodeMap {
 
 class Node extends ClassicPreset.Node {
   width = 180;
-  height = 120;
+  height = 100;
 }
 class Connection<N extends Node> extends ClassicPreset.Connection<N, N> {}
 
@@ -57,27 +57,30 @@ type AreaExtra = ReactArea2D<any> | MinimapExtra | ContextMenuExtra;
 
 export function createDynamicNode(nodeData: any) {
   const node = new Node(nodeData.label);
-
+  // resize the node based on the max length of the input/output names
+  let maxSocketNameLength = 0;
   nodeData.inputs.forEach((input: NodeInput) => {
     let socket = new ClassicPreset.Socket(input.name);
     if (!node.inputs.hasOwnProperty(input.name)) {
       node.addInput(input.name, new ClassicPreset.Input(socket, input.name));
-      node.height += 25; // Increase height of node for each input
+      maxSocketNameLength = Math.max(maxSocketNameLength, input.name.length);
     }
-});
+  });
 
   nodeData.outputs.forEach((output: NodeOutput) => {
     let socket = new ClassicPreset.Socket(output.name);
     if (!node.outputs.hasOwnProperty(output.name)) {
       node.addOutput(output.name, new ClassicPreset.Output(socket, output.name));
-      node.height += 25; // Increase height of node for each output
-  }
+      maxSocketNameLength = Math.max(maxSocketNameLength, output.name.length);
+    }
   });
+  node.height = Math.max(140, node.height + (nodeData.inputs.length + nodeData.outputs.length) * 35)
+  node.width += maxSocketNameLength * 5;
 
   return node;
 }
 
-export async function addNode(editor, nodeData) {
+export async function addNode(editor, area, nodeData) {
   console.log("Adding node", nodeData);
   const node = createDynamicNode(nodeData);
   await editor.addNode(node);
@@ -138,7 +141,7 @@ export async function removeNode(editor, name) {
   });
 }
 
-export async function createEditor(container: HTMLElement, worktreeData: any) {
+export async function createEditor(container: HTMLElement, settings: any, worktreeData: any) {
   container.innerHTML = ''
 
   const editor = new NodeEditor<Schemes>();
@@ -179,7 +182,9 @@ export async function createEditor(container: HTMLElement, worktreeData: any) {
   area.use(render);
   area.use(arrange);
   area.use(contextMenu);
-  area.use(minimap);
+  if (settings.minimap) {
+    area.use(minimap);
+  }
 
   AreaExtensions.simpleNodesOrder(area);
 
@@ -193,9 +198,11 @@ export async function createEditor(container: HTMLElement, worktreeData: any) {
   const nodeMap: NodeMap = {}; // To keep track of created nodes for linking
   editor.nodeMap = nodeMap;
 
+  console.log("worktreeData", worktreeData)
+  console.log("settings: ", settings)
   for (const nodeId in worktreeData.nodes) {
     const nodeData = worktreeData.nodes[nodeId];
-    await addNode(editor, nodeData);
+    await addNode(editor, area, nodeData);
   }
 
   // Adding connections based on worktreeData

aiida_worktree/widget/js/widget.css

@@ -24,8 +24,6 @@
 
 .rete {
 	position: relative;
-	width: 90%;
-	height: 600px;
 	font-size: 1rem;
 	background: rgb(200, 190, 190);
 	margin: 1em;

aiida_worktree/widget/js/widget.tsx

@@ -4,8 +4,8 @@ import { createEditor, removeNode, addNode, addLink, removeLink } from "./defaul
 import "./widget.css";
 
 export function useRete<T extends { destroy(): void }>(
-  create: (el: HTMLElement, data: any) => Promise<T>,
-  worktreeData: any
+  create: (el: HTMLElement, settings: any, data: any) => Promise<T>,
+  settings: any, worktreeData: any
 ) {
   const [container, setContainer] = React.useState<null | HTMLElement>(null);
   const editorRef = React.useRef<T>(); // Ref for storing the actual editor instance
@@ -19,7 +19,7 @@ export function useRete<T extends { destroy(): void }>(
         editorRef.current.destroy();
         container.innerHTML = '';
       }
-      create(container, worktreeData).then((value) => {
+      create(container, settings, worktreeData).then((value) => {
         editorRef.current = value;
         setEditor(value);
         // Expose the editor instance to the window object for debugging
@@ -51,8 +51,10 @@ export function useRete<T extends { destroy(): void }>(
 }
 
 const render = createRender(() => {
-  const [value, setValue] = useModelState<number>("value");
-  const [ref, editor] = useRete(createEditor, value);
+  const [value, setValue] = useModelState<any>("value");
+  const [style, setStyle] = useModelState<any>("style");
+  const [settings, setSettings] = useModelState<any>("settings");
+  const [ref, editor] = useRete(createEditor, settings, value);
   const model = useModel();
 
   React.useEffect(() => {
@@ -93,7 +95,7 @@ const render = createRender(() => {
 
   return (
     <div className="App">
-      <div ref={ref} className="rete"></div>
+      <div ref={ref} className="rete" style={style}></div>
     </div>
   );
 });

aiida_worktree/widget/src/widget/__init__.py

@@ -11,16 +11,28 @@
     __version__ = "unknown"
 
 default_value = {"summary": {}, "nodes": {}, "links": [], "logs": [], "pk": []}
+default_style = {"width": "90%", "height": "600px"}
+default_settings = {"minimap": True}
 
 
 class NodeGraphWidget(anywidget.AnyWidget):
     _esm = pathlib.Path(__file__).parent / "static" / "widget.js"
     _css = pathlib.Path(__file__).parent / "static" / "widget.css"
     value = traitlets.Dict(default_value).tag(sync=True)
+    settings = traitlets.Dict(default_settings).tag(sync=True)
+    style = traitlets.Dict(default_style).tag(sync=True)
 
     def from_worktree(self, worktree):
         from aiida_worktree.web.backend.app.utils import worktree_to_short_json
 
         wtdata = worktree.to_dict()
         wait_to_link(wtdata)
-        self.value = worktree_to_short_json(wtdata)
+        wtdata = worktree_to_short_json(wtdata)
+        self.value = wtdata
+
+    def from_node(self, node):
+        ndata = node.to_dict()
+        ndata.pop("properties", None)
+        ndata["label"] = ndata["metadata"]["identifier"]
+        wtdata = {"nodes": {node.name: ndata}, "links": []}
+        self.value = wtdata

docs/source/howto/if.ipynb

@@ -586,7 +586,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.0"
+   "version": "3.11.0"
   },
   "vscode": {
    "interpreter": {