Add entity group -column to the "Add entities" -dialog

The column is optional to fill. If filled, the entity group will be
created if it doesn't already exist, and the entity will be added
to the group.

#2788

spinetoolbox/spine_db_editor/widgets/add_items_dialogs.py

@@ -440,8 +440,10 @@ def __init__(self, parent, item, db_mngr, *db_maps, force_default=False, commit_
             self.entity_names_by_class_name = {entity_class_name: entity_name}
             if item.parent_item.item_type == "entity_class":
                 self.class_key = item.parent_item.display_id
+                self.class_item = item.parent_item
         elif item.item_type == "entity_class":
             self.class_key = item.display_id
+            self.class_item = item
         self.model.force_default = force_default
         self.setWindowTitle("Add entities")
         self.table_view.setItemDelegate(ManageEntitiesDelegate(self))
@@ -484,25 +486,45 @@ def make_model(self):
         return EmptyRowModel(self)
 
     def _do_reset_model(self):
-        header = self.dimension_name_list + ("entity name", "alternative", "databases")
+        header = self.dimension_name_list + ("entity name", "alternative", "entity group", "databases")
         self.model.set_horizontal_header_labels(header)
         default_db_maps = [db_map for db_map, keys in self.db_map_ent_cls_lookup.items() if self.class_key in keys]
         db_names = ",".join([db_name for db_name, db_map in self.keyed_db_maps.items() if db_map in default_db_maps])
         alt_selection_model = self.parent().ui.alternative_tree_view.selectionModel()
         alt_selection = alt_selection_model.selection()
-        alternative = None
+        selected_alt_name = None
         if alt_selection:
-            selected_alternative_index = alt_selection_model.currentIndex()
-            alternative = selected_alternative_index.model().item_from_index(selected_alternative_index)
+            selected_alt_index = alt_selection_model.currentIndex()
+            selected_alt = selected_alt_index.model().item_from_index(selected_alt_index)
+            selected_alt_name = self.append_db_codenames(selected_alt.name, {selected_alt.db_map})
         all_databases = [db_map for db_name, db_map in self.keyed_db_maps.items() if db_map in default_db_maps]
-        alt_name_list = [x["name"] for db_map in all_databases for x in self.db_mngr.get_items(db_map, "alternative")]
-        alt_name_list.append("")
-        alternative_name = alternative.name if alternative else alt_name_list[0]
+        dbs_by_alternative = {}
+        for db_map in all_databases:
+            for alternative in self.db_mngr.get_items(db_map, "alternative"):
+                dbs_by_alternative.setdefault(alternative["name"], set()).add(db_map)
+        alt_names_list = []
+        for alternative, db_maps in dbs_by_alternative.items():
+            alt_names_list.append(self.append_db_codenames(alternative, db_maps))
+        alt_names_list.append("")  # If the db has no alternatives, defaults to nothing
+        alternative_name = selected_alt_name if selected_alt_name else alt_names_list[0]
         defaults = {"databases": db_names, "alternative": alternative_name}
         defaults.update(self.entity_names_by_class_name)
         self.model.set_default_row(**defaults)
         self.model.clear()
 
+    def append_db_codenames(self, name, db_maps):
+        """Appends a list of given databases to the given name
+
+        Args:
+            name (str): The name where the database names should be appended
+            db_maps (set): A set containing DatabaseMapping objects
+        Returns:
+            str: The name with the databases appended to the end after an @ sign.
+        """
+        if len(db_maps) == len(self.parent().db_maps):
+            return name
+        return name + "@(" + ", ".join(db_map.codename for db_map in db_maps) + ")"
+
     def get_db_map_data(self):
         db_map_data = {}
         name_column = self.model.horizontal_header_labels().index("entity name")
@@ -564,6 +586,9 @@ def accept(self):
         entity_alternatives = self.make_entity_alternatives(created_entities)
         if entity_alternatives:  # If alternatives have been defined
             self.db_mngr.add_entity_alternatives(entity_alternatives)
+        entity_groups = self.make_entity_groups(created_entities)
+        if entity_groups:  # If entity groups have been defined
+            self.db_mngr.import_data(entity_groups, command_text="Add entity group")
         super().accept()
 
     def make_entity_alternatives(self, entities):
@@ -577,6 +602,7 @@ def make_entity_alternatives(self, entities):
             alternative = row_data[alternative_column]
             if not alternative:
                 continue
+            alternative = alternative.split("@")[0]
             entity_name = row_data[name_column]
             entity = entities[entity_name]
             db_names = row_data[db_column]
@@ -592,6 +618,30 @@ def make_entity_alternatives(self, entities):
                 )
         return entity_alternatives
 
+    def make_entity_groups(self, entities):
+        """Creates a mapping from db_map to entity alternatives that are to be created"""
+        name_column = self.model.horizontal_header_labels().index("entity name")
+        entity_group_column = self.model.horizontal_header_labels().index("entity group")
+        db_column = self.model.horizontal_header_labels().index("databases")
+        db_map_data = {}
+        for i in range(self.model.rowCount() - 1):
+            row_data = self.model.row_data(i)
+            entity_group = row_data[entity_group_column]
+            if not entity_group:
+                continue
+            entity_group = entity_group.split("@")[0]
+            entity_name = row_data[name_column]
+            entity = entities[entity_name]
+            class_name = entity["entity_class_name"]
+            db_names = row_data[db_column]
+            for db_name in db_names.split(","):
+                db_map = self.keyed_db_maps[db_name]
+                db_map_data.setdefault(db_map, {}).setdefault("entities", set()).add((class_name, entity_group))
+                db_map_data.setdefault(db_map, {}).setdefault("entity_groups", set()).add(
+                    (class_name, entity_group, entity["name"])
+                )
+        return db_map_data
+
 
 class ManageElementsDialog(AddEntitiesOrManageElementsDialog):
     """A dialog to query user's preferences for managing entity dimensions."""

spinetoolbox/spine_db_editor/widgets/custom_delegates.py

@@ -722,10 +722,51 @@ def _create_alternative_editor(self, parent, index):
             QWidget: editor
         """
         editor = SearchBarEditor(self.parent(), parent)
-        all_databases = self.parent().keyed_db_maps.values()
-        name_list = [
-            x["name"] for db_map in all_databases for x in self.parent().db_mngr.get_items(db_map, "alternative")
-        ]
+        name_list = set()
+        dbs_by_alternative_name = {}
+        database_column = self.parent().model.horizontal_header_labels().index("databases")
+        database_index = index.model().index(index.row(), database_column)
+        databases = database_index.data(Qt.ItemDataRole.DisplayRole).split(",")
+        for db_map_codename in databases:  # Filter possible alternatives based on selected databases
+            db_map = self.parent().keyed_db_maps[db_map_codename]
+            alternatives = self.parent().db_mngr.get_items(db_map, "alternative")
+            for alternative in alternatives:
+                dbs_by_alternative_name.setdefault(alternative["name"], set()).add(db_map)
+        for alternative, db_maps in dbs_by_alternative_name.items():
+            name_list.add(self.parent().append_db_codenames(alternative, db_maps))
+        editor.set_data(index.data(Qt.ItemDataRole.EditRole), name_list)
+        return editor
+
+    def _create_entity_group_editor(self, parent, index):
+        """Creates an editor.
+
+        Args:
+            parent (QWidget): parent widget
+            index (QModelIndex): index being edited
+
+        Returns:
+            QWidget: editor
+        """
+        database_column = self.parent().model.horizontal_header_labels().index("databases")
+        database_index = index.model().index(index.row(), database_column)
+        databases = database_index.data(Qt.ItemDataRole.DisplayRole).split(",")
+        entity_class = self.parent().class_item
+        dbs_by_entity_group = {}  # A mapping from entity_group to db_map(s)
+        for db_map in entity_class.db_maps:
+            if db_map.codename not in databases:  # Allow groups that are in selected DBs under "databases" -column.
+                continue
+            class_item = self.parent().db_mngr.get_item_by_field(db_map, "entity_class", "name", entity_class.name)
+            if not class_item:
+                continue
+            ent_groups = self.parent().db_mngr.get_items_by_field(
+                db_map, "entity_group", "entity_class_id", class_item["id"]
+            )
+            for ent_group in ent_groups:
+                dbs_by_entity_group.setdefault(ent_group["group_name"], set()).add(db_map)
+        name_list = set()
+        for entity_group, db_maps in dbs_by_entity_group.items():
+            name_list.add(self.parent().append_db_codenames(entity_group, db_maps))
+        editor = SearchBarEditorWithCreation(self.parent(), parent)
         editor.set_data(index.data(Qt.ItemDataRole.EditRole), name_list)
         return editor
 
@@ -807,6 +848,8 @@ def createEditor(self, parent, option, index):
             editor = self._create_database_editor(parent, index)
         elif header[index.column()] == "alternative":
             editor = self._create_alternative_editor(parent, index)
+        elif header[index.column()] == "entity group":
+            editor = self._create_entity_group_editor(parent, index)
         else:
             editor = SearchBarEditor(parent)
             entity_name_list = self.parent().entity_name_list(index.row(), index.column())

tests/spine_db_editor/widgets/helpers.py

@@ -109,17 +109,14 @@ def add_zero_dimension_entity_class(view, name):
     add_entity_tree_item({0: name}, view, "Add entity classes", AddEntityClassesDialog)
 
 
-def add_entity(view, name, entity_class_index=0):
+def add_entity(view, name, entity_class_index=0, alternative=None, group=None):
     model = view.model()
     root_index = model.index(0, 0)
     class_index = model.index(entity_class_index, 0, root_index)
     view._context_item = model.item_from_index(class_index)
-    add_entity_tree_item({0: name}, view, "Add entities", AddEntitiesDialog)
-
-
-def add_entity_with_alternative(view, name, alternative, entity_class_index=0):
-    model = view.model()
-    root_index = model.index(0, 0)
-    class_index = model.index(entity_class_index, 0, root_index)
-    view._context_item = model.item_from_index(class_index)
-    add_entity_tree_item({0: name, 1: alternative}, view, "Add entities", AddEntitiesDialog)
+    data = {0: name}
+    if alternative:
+        data.update({1: alternative})
+    if group:
+        data.update({2: group})
+    add_entity_tree_item(data, view, "Add entities", AddEntitiesDialog)

tests/spine_db_editor/widgets/test_custom_qtreeview.py

@@ -36,7 +36,6 @@
     add_entity_tree_item,
     add_zero_dimension_entity_class,
     add_entity,
-    add_entity_with_alternative,
 )
 
 
@@ -196,10 +195,12 @@ def test_add_entity(self):
         self.assertEqual(data[0].name, "an_entity")
 
     def test_add_entity_with_alternative(self):
+        """Tests that adding a new entity with the alternative -column filled
+        will actually add the alternative"""
         view = self._db_editor.ui.treeView_entity
         add_zero_dimension_entity_class(view, "an_entity_class")
         with mock.patch.object(self._db_mngr, "add_entity_alternatives") as mock_add_entity_alternatives:
-            add_entity_with_alternative(view, "an_entity", "Alt1")
+            add_entity(view, "an_entity", alternative="Alt1")
             mock_add_entity_alternatives.assert_called_once_with(
                 {
                     self._db_map: [
@@ -233,6 +234,43 @@ def test_add_entity_with_alternative(self):
         self.assertEqual(len(data), 1)
         self.assertEqual(data[0].name, "an_entity")
 
+    def test_add_entity_with_group(self):
+        """Tests that adding a new entity with the group -column filled
+        will actually create the entity group and add the entity to it."""
+        view = self._db_editor.ui.treeView_entity
+        add_zero_dimension_entity_class(view, "classy")
+        with mock.patch.object(self._db_mngr, "import_data") as mock_add_entity_groups:
+            add_entity(view, "wine", group="beverages")
+            mock_add_entity_groups.assert_called_once_with(
+                {
+                    self._db_map: {
+                        "entities": {("classy", "beverages")},
+                        "entity_groups": {("classy", "beverages", "wine")},
+                    }
+                },
+                command_text="Add entity group",
+            )
+        model = view.model()
+        root_index = model.index(0, 0)
+        class_index = model.index(0, 0, root_index)
+        model.fetchMore(class_index)
+        while model.rowCount(class_index) != 1:
+            QApplication.processEvents()
+        self.assertEqual(model.rowCount(class_index), 1)
+        self.assertEqual(class_index.data(), "classy")
+        entity_index = model.index(0, 0, class_index)
+        self.assertEqual(model.rowCount(entity_index), 0)
+        self.assertEqual(entity_index.data(), "wine")
+        entity_database_index = model.index(0, 1, class_index)
+        self.assertEqual(entity_database_index.data(), self.db_codename)
+        self._commit_changes_to_database("Add entity.")
+        data = self._db_map.query(self._db_map.entity_class_sq).all()
+        self.assertEqual(len(data), 1)
+        self.assertEqual(data[0].name, "classy")
+        data = self._db_map.query(self._db_map.entity_sq).all()
+        self.assertEqual(len(data), 1)
+        self.assertEqual(data[0].name, "wine")
+
     def test_add_entity_with_single_dimension(self):
         view = self._db_editor.ui.treeView_entity
         add_zero_dimension_entity_class(view, "an_entity_class")