🪲 Fix socket missing when converting from .ipynb

pyflow/scene/from_ipynb_conversion.py

@@ -169,22 +169,27 @@ def get_edges_data(blocks_data: OrderedDict) -> OrderedDict:
     if len(blocks_data) > 0:
         greatest_block_id = blocks_data[-1]["id"]
 
-    for i in range(1, len(code_blocks)):
-        socket_id_out = greatest_block_id + 2 * i + 2
-        socket_id_in = greatest_block_id + 2 * i + 1
-        code_blocks[i - 1]["sockets"].append(
-            get_output_socket_data(socket_id_out, code_blocks[i - 1]["width"])
-        )
-        code_blocks[i]["sockets"].append(get_input_socket_data(socket_id_in))
-        edges_data.append(
-            get_edge_data(
-                i,
-                code_blocks[i - 1]["id"],
-                socket_id_out,
-                code_blocks[i]["id"],
-                socket_id_in,
+    last_socket_id_out: int = -1
+    for i, block in enumerate(code_blocks):
+        socket_id_out: int = greatest_block_id + 2 * i + 2
+        socket_id_in: int = greatest_block_id + 2 * i + 1
+
+        block["sockets"].append(get_output_socket_data(socket_id_out, block["width"]))
+        block["sockets"].append(get_input_socket_data(socket_id_in))
+
+        if i >= 1:
+            edges_data.append(
+                get_edge_data(
+                    i,
+                    code_blocks[i - 1]["id"],
+                    last_socket_id_out,
+                    code_blocks[i]["id"],
+                    socket_id_in,
+                )
             )
-        )
+
+        last_socket_id_out = socket_id_out
+
     return edges_data
 
 

tests/unit/scene/test_ipynb_conversion.py

@@ -3,7 +3,7 @@
 
 """Unit tests for the conversion from and to ipynb."""
 
-from typing import OrderedDict
+from typing import List, OrderedDict
 from pytest_mock import MockerFixture
 import pytest_check as check
 import json
@@ -71,7 +71,7 @@ def check_conversion_coherence(ipynb_data: OrderedDict, ipyg_data: OrderedDict):
     2. the right amount of code blocks and edges
     3. blocks and sockets with unique ids
     4. edges with existing ids
-    5. code blocks that always have a source
+    5. code blocks that always have a source and two sockets
     6. markdown blocks that always have text
     """
 
@@ -109,10 +109,17 @@ def check_conversion_coherence(ipynb_data: OrderedDict, ipyg_data: OrderedDict):
         check.equal(edge["source"]["socket"] in socket_id_set, True)
         check.equal(edge["destination"]["socket"] in socket_id_set, True)
 
-    # code blocks always have a source and markdown blocks always have a text
+    # code blocks always have a source and two sockets
+    # markdown blocks always have a text
     for block in ipyg_data["blocks"]:
         if block["block_type"] == BLOCK_TYPE_TO_NAME["code"]:
-            check.equal("source" in block and type(block["source"]) == str, True)
+            check.equal("source" in block, True)
+            check.equal(type(block["source"]), str)
+
+            check.equal("sockets" in block, True)
+            check.equal(type(block["sockets"]), list)
+            check.equal(len(block["sockets"]), 2)
+
         if block["block_type"] == BLOCK_TYPE_TO_NAME["markdown"]:
             check.equal("text" in block and type(block["text"]) == str, True)