fix(serializer): Fix Enumeration labels

capellambse_context_diagrams/serializers.py

@@ -136,7 +136,7 @@ class type that stores all previously named classes.
             size = (child["size"]["width"], child["size"]["height"])  # type: ignore
             features = []
             if styleclass in decorations.needs_feature_line:
-                features = handle_features(child)  # type: ignore[arg-type]
+                features = handle_features(child, styleclass)  # type: ignore[arg-type]
 
             element = diagram.Box(
                 ref,
@@ -302,23 +302,17 @@ def order_children(self) -> None:
         self.diagram = new_diagram
 
 
-def handle_features(child: _elkjs.ELKOutputNode) -> list[str]:
+def handle_features(child: _elkjs.ELKOutputNode, styleclass: str) -> list[str]:
     """Return all consecutive labels (without first) from the ``child``."""
     features: list[str] = []
     if len(child["children"]) <= 1:
         return features
 
-    labels: cabc.MutableSequence[_elkjs.ELKOutputChild] = []
-    for c in child["children"]:
-        if c["type"] != "label":
-            continue
-
-        if ":" not in c["text"]:
-            labels.append(c)
-            continue
-
-        features.append(c["text"])
-    child["children"] = labels
+    feature_collector = FEATURE_HANDLER.get(
+        styleclass, collect_class_properties
+    )
+    labels, features = feature_collector(child["children"])
+    child["children"] = labels  # type: ignore[typeddict-item]
     return features
 
 
@@ -351,6 +345,55 @@ def reverse_edge_refpoints(child: _elkjs.ELKOutputEdge) -> None:
     child["routingPoints"] = child["routingPoints"][::-1]
 
 
+def collect_class_properties(
+    children: cabc.Iterable[_elkjs.ELKOutputChild],
+) -> tuple[list[_elkjs.ELKOutputLabel], list[str]]:
+    features: list[str] = []
+    labels: list[_elkjs.ELKOutputLabel] = []
+    for child in children:
+        if child["type"] != "label":
+            continue
+
+        if ":" not in child["text"]:
+            labels.append(child)
+            continue
+
+        features.append(child["text"])
+    return labels, features
+
+
+def collect_enumeration_properties(
+    children: cabc.MutableSequence[_elkjs.ELKOutputChild],
+) -> tuple[list[_elkjs.ELKOutputLabel], list[str]]:
+    features: list[str] = []
+    labels: list[_elkjs.ELKOutputLabel] = []
+    for i, child in enumerate(children):
+        if child["type"] != "label":
+            continue
+
+        if (
+            child["text"] == child["text"].upper()
+            and i > 0
+            or ":" in child["text"]
+        ):
+            features.append(child["text"])
+            continue
+
+        labels.append(child)
+    return labels, features
+
+
 EDGE_HANDLER: dict[str, cabc.Callable[[_elkjs.ELKOutputEdge], None]] = {
     "Generalization": reverse_edge_refpoints
 }
+LabelsAndFeatures = tuple[list[_elkjs.ELKOutputLabel], list[str]]
+FEATURE_HANDLER: dict[
+    str,
+    cabc.Callable[
+        [cabc.MutableSequence[_elkjs.ELKOutputChild]], LabelsAndFeatures
+    ],
+] = {
+    "Class": collect_class_properties,
+    "PrimitiveClass": collect_class_properties,
+    "Enumeration": collect_enumeration_properties,
+}