From b75461db442fccad92ff7ce47f43e15b0c30b76e Mon Sep 17 00:00:00 2001
From: huyenngn <ngan-huyen.thi-nguyen@deutschebahn.com>
Date: Tue, 20 Aug 2024 19:37:20 +0200
Subject: [PATCH] fix: Fix port spread calculation

---
 .../collectors/default.py                     | 30 +++++++++++++------
 .../collectors/generic.py                     |  2 +-
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/capellambse_context_diagrams/collectors/default.py b/capellambse_context_diagrams/collectors/default.py
index 71a31732..1d2eb37e 100644
--- a/capellambse_context_diagrams/collectors/default.py
+++ b/capellambse_context_diagrams/collectors/default.py
@@ -109,22 +109,32 @@ def process_context(self):
             self.centerbox.height, *stack_heights.values()
         )
 
+    def _process_port_spread(self, exs, attr, port_spread, owners) -> None:
+        for ex in exs:
+            elem = getattr(ex, attr).owner
+            if (owner := owners.get(elem.uuid)) is None:
+                owner = [
+                    uuid
+                    for uuid in generic.get_all_owners(elem)
+                    if uuid not in self.diagram_target_owners
+                ][-1]
+                owners[elem.uuid] = owner
+            port_spread.setdefault(owner, 0)
+            port_spread[owner] += 1
+
     def _process_exchanges(self) -> tuple[
         list[common.GenericElement],
         list[generic.ExchangeData],
     ]:
         inc, out = port_collector(self.diagram.target, self.diagram.type)
-        port_spread: dict[str, int] = {}
         inc_c = port_exchange_collector(inc)
         out_c = port_exchange_collector(out)
         inc_exchanges = list(chain.from_iterable(inc_c.values()))
         out_exchanges = list(chain.from_iterable(out_c.values()))
-        for ex in inc_exchanges:
-            port_spread.setdefault(ex.source.owner.uuid, 0)
-            port_spread[ex.source.owner.uuid] += 1
-        for ex in out_exchanges:
-            port_spread.setdefault(ex.target.owner.uuid, 0)
-            port_spread[ex.target.owner.uuid] -= 1
+        port_spread: dict[str, int] = {}
+        owners: dict[str, str] = {}
+        self._process_port_spread(inc_exchanges, "source", port_spread, owners)
+        self._process_port_spread(out_exchanges, "target", port_spread, owners)
         self.exchanges = inc_exchanges + out_exchanges
         ex_datas: list[generic.ExchangeData] = []
         for ex in self.exchanges:
@@ -148,12 +158,14 @@ def _process_exchanges(self) -> tuple[
                     is_hierarchical,
                 )
                 src, tgt = generic.exchange_data_collector(ex_data)
+                src_owner = owners.get(src.owner.uuid, "-1")
+                tgt_owner = owners.get(tgt.owner.uuid, "-1")
                 if (
                     (src.parent == self.diagram.target)
-                    and (port_spread.get(tgt.parent.uuid, 0) > 0)
+                    and (port_spread.get(tgt_owner, 0) > 0)
                 ) or (
                     (tgt.parent == self.diagram.target)
-                    and (port_spread.get(src.parent.uuid, 0) <= 0)
+                    and (port_spread.get(src_owner, 0) <= 0)
                 ):
                     elkdata.edges[-1].sources = [tgt.uuid]
                     elkdata.edges[-1].targets = [src.uuid]
diff --git a/capellambse_context_diagrams/collectors/generic.py b/capellambse_context_diagrams/collectors/generic.py
index 41927d9e..94c8e6ea 100644
--- a/capellambse_context_diagrams/collectors/generic.py
+++ b/capellambse_context_diagrams/collectors/generic.py
@@ -254,7 +254,7 @@ def move_edges(
 
 def get_all_owners(obj: common.GenericElement) -> cabc.Iterator[str]:
     """Return the UUIDs from all owners of ``obj``."""
-    current = obj
+    current: common.GenericElement | None = obj
     while current is not None:
         yield current.uuid
         current = getattr(current, "owner", None)