Feat custom diagram

capellambse_context_diagrams/__init__.py

@@ -65,10 +65,12 @@ def init() -> None:
     """Initialize the extension."""
     register_classes()
     register_interface_context()
+    register_physical_port_context()
     register_tree_view()
     register_realization_view()
     register_data_flow_view()
     register_cable_tree_view()
+    register_custom_diagram()
     # register_functional_context() XXX: Future
 
 
@@ -247,6 +249,15 @@ def register_functional_context() -> None:
         )
 
 
+def register_physical_port_context() -> None:
+    """Add the `context_diagram` attribute to `PhysicalPort`s."""
+    m.set_accessor(
+        cs.PhysicalPort,
+        ATTR_NAME,
+        context.PhysicalPortContextAccessor(DiagramType.PAB.value, {}),
+    )
+
+
 def register_tree_view() -> None:
     """Add the ``tree_view`` attribute to ``Class``es."""
     m.set_accessor(
@@ -313,3 +324,31 @@ def register_cable_tree_view() -> None:
             {},
         ),
     )
+
+
+def register_custom_diagram() -> None:
+    """Add the `custom_diagram` attribute to `ModelObject`s."""
+    supported_classes: list[tuple[type[m.ModelElement], DiagramType]] = [
+        (oa.Entity, DiagramType.OAB),
+        (oa.OperationalActivity, DiagramType.OAB),
+        (oa.OperationalCapability, DiagramType.OCB),
+        (oa.CommunicationMean, DiagramType.OAB),
+        (sa.Mission, DiagramType.MCB),
+        (sa.Capability, DiagramType.MCB),
+        (sa.SystemComponent, DiagramType.SAB),
+        (sa.SystemFunction, DiagramType.SAB),
+        (la.LogicalComponent, DiagramType.LAB),
+        (la.LogicalFunction, DiagramType.LAB),
+        (pa.PhysicalComponent, DiagramType.PAB),
+        (pa.PhysicalFunction, DiagramType.PAB),
+        (cs.PhysicalLink, DiagramType.PAB),
+        (cs.PhysicalPort, DiagramType.PAB),
+        (fa.ComponentExchange, DiagramType.SAB),
+        (information.Class, DiagramType.CDB),
+    ]
+    for class_, dgcls in supported_classes:
+        m.set_accessor(
+            class_,
+            "custom_diagram",
+            context.CustomContextAccessor(dgcls.value, {}),
+        )