From 5ee68c614ad9d0dc7b41cf16870414241731d011 Mon Sep 17 00:00:00 2001
From: huyenngn <ngan-huyen.thi-nguyen@deutschebahn.com>
Date: Mon, 18 Dec 2023 10:24:54 +0100
Subject: [PATCH] feat: Add context to elements

---
 .vscode/launch.json                     | 15 +++++++++++++++
 capellambse_context_diagrams/context.py | 14 ++++++++------
 2 files changed, 23 insertions(+), 6 deletions(-)
 create mode 100644 .vscode/launch.json

diff --git a/.vscode/launch.json b/.vscode/launch.json
new file mode 100644
index 00000000..25882f8b
--- /dev/null
+++ b/.vscode/launch.json
@@ -0,0 +1,15 @@
+{
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Test ClassTree Diagram",
+            "type": "python",
+            "request": "launch",
+            "console": "integratedTerminal",
+            "envFile": ".env",
+            "pythonArgs": ["-Xdev"],
+            "module": "capellambse.repl",
+            "args": ["ros-msg"]
+        }
+    ]
+}
diff --git a/capellambse_context_diagrams/context.py b/capellambse_context_diagrams/context.py
index 9635d0c5..28be2b1a 100644
--- a/capellambse_context_diagrams/context.py
+++ b/capellambse_context_diagrams/context.py
@@ -364,8 +364,9 @@ def _create_diagram(self, params: dict[str, t.Any]) -> cdiagram.Diagram:
         )
         data, legend = tree_view.collector(self, params)
         params["elkdata"] = data
-        add_context(data)
-        class_diagram = super()._create_diagram(params)
+        layout = _elkjs.call_elkjs(data)
+        add_context(layout)
+        class_diagram = self.serializer.make_diagram(layout)
         width, height = class_diagram.viewport.size
         axis: t.Literal["x", "y"]
         if params["elk.direction"] in {"DOWN", "UP"}:
@@ -380,18 +381,19 @@ def _create_diagram(self, params: dict[str, t.Any]) -> cdiagram.Diagram:
         return class_diagram
 
 
-def add_context(data: _elkjs.ELKInputData) -> None:
+def add_context(data: _elkjs.ELKOutputData) -> None:
     ids: set[str] = set()
 
-    def get_ids(obj: _elkjs.ELKInputChild):
-        ids.add(obj["id"])
+    def get_ids(obj: _elkjs.ELKOutputChild):
+        if obj["id"] and not obj["id"].startswith("g_"):
+            ids.add(obj["id"])
         for cobj in obj.get("children", []):
             get_ids(cobj)
 
     for child in data["children"]:
         get_ids(child)
 
-    for child in data["children"] + data["edges"]:
+    for child in data["children"]:
         child["context"] = list(ids)