From c3d890cf8d81dc44095c386b87d9515deb0d5bb2 Mon Sep 17 00:00:00 2001
From: Martin Lehmann <martin.lehmann@deutschebahn.com>
Date: Thu, 19 Dec 2024 11:44:14 +0100
Subject: [PATCH] feat: Improve class selection during insert without type hint

---
 capellambse/model/_descriptors.py | 60 ++++++++++++++++++++++++++++---
 1 file changed, 56 insertions(+), 4 deletions(-)

diff --git a/capellambse/model/_descriptors.py b/capellambse/model/_descriptors.py
index 06cec4e3..ee4aab25 100644
--- a/capellambse/model/_descriptors.py
+++ b/capellambse/model/_descriptors.py
@@ -51,7 +51,7 @@
 import capellambse
 from capellambse import helpers
 
-from . import T, T_co, U_co
+from . import T, T_co, U, U_co
 
 _NotSpecifiedType = t.NewType("_NotSpecifiedType", object)
 _NOT_SPECIFIED = _NotSpecifiedType(object())
@@ -466,6 +466,8 @@ def insert(
         elmlist: _obj.ElementListCouplingMixin,
         index: int,
         value: T_co | NewObject,
+        *,
+        bounds: tuple[_obj.ClassName, ...] = (),
     ) -> T_co:
         """Insert the ``value`` object into the model.
 
@@ -2855,6 +2857,8 @@ def insert(
         elmlist: _obj.ElementListCouplingMixin,
         index: int,
         value: T_co | NewObject,
+        *,
+        bounds: tuple[_obj.ClassName, ...] = (),
     ) -> T_co:
         if self.role_tag is None:
             raise RuntimeError(
@@ -2885,9 +2889,18 @@ def insert(
 
         if isinstance(value, NewObject):
             if not value._type_hint:
-                cls = self._guess_xtype(elmlist._model)
+                cls = self._guess_xtype(elmlist._model, bounds)
             else:
                 cls = self._match_xtype(elmlist._model, value._type_hint)
+
+                for bound in bounds:
+                    bound_cls = elmlist._model.resolve_class(bound)
+                    if not isinstance(cls, bound_cls):
+                        raise InvalidModificationError(
+                            f"Requested class {cls!r}"
+                            f" does not satisfy bound {bound}"
+                        )
+
             attrs = dict(value._kw)
             if not hasattr(cls, "__capella_namespace__"):
                 raise TypeError(
@@ -2960,8 +2973,40 @@ def _match_xtype(
             raise TypeError(f"Cannot insert elements into {self._qualname!r}")
         return t.cast(type[T_co], cls)
 
-    def _guess_xtype(self, model: capellambse.MelodyModel) -> type[T_co]:
-        return t.cast(type[T_co], model.resolve_class(self.class_))
+    def _guess_xtype(
+        self,
+        model: capellambse.MelodyModel,
+        bounds: tuple[_obj.ClassName, ...],
+    ) -> type[T_co]:
+        cls = t.cast(type[T_co], model.resolve_class(self.class_))
+        if not bounds:
+            return cls
+
+        candidates = _find_all_subclasses(cls)
+
+        for bound in bounds:
+            candidates &= _find_all_subclasses(model.resolve_class(bound))
+
+        if not candidates:
+            strbounds = ", ".join(f"{b[0].alias}:{b[1]}" for b in bounds)
+            raise InvalidModificationError(
+                f"No subclass of {self.class_[0].alias}:{self.class_[1]}"
+                f" satisfies all bounds: {strbounds}"
+            )
+
+        filtered_candidates: set[type[T_co]] = set()
+        for c in candidates:
+            if c.__capella_abstract__ or any(
+                c != o and issubclass(c, o) for o in candidates
+            ):
+                filtered_candidates.add(c)
+        candidates -= filtered_candidates
+        assert candidates
+
+        if len(candidates) > 1:
+            raise RuntimeError("Multiple eligible classes found: {candidates}")
+
+        return candidates.pop()
 
     def _resolve_super_attributes(
         self, super_acc: Accessor[t.Any] | None
@@ -3014,6 +3059,13 @@ def make_coupled_list_type(
     return list_type
 
 
+def _find_all_subclasses(cls: type[U]) -> set[type[U]]:
+    classes = set(cls.__subclasses__())
+    for scls in classes.copy():
+        classes.update(_find_all_subclasses(scls))
+    return classes
+
+
 from . import _obj
 from ._obj import Namespace  # noqa: F401 # needed for Sphinx