Remove _add_images_to_existing_mixed_dataset and _add_recon_to_dataset_and_tree_view

Author: samtygier-stfc

mantidimaging/core/data/dataset.py

@@ -178,6 +178,16 @@ def set_stack(self, file_type: FILE_TYPES, image_stack: ImageStack) -> None:
         else:
             raise AttributeError(f"StrictDataset does not have an attribute for {attr_name}")
 
+    def set_stack_by_type_name(self, file_type_name: str, image_stack: ImageStack) -> None:
+        file_type_name = file_type_name.upper().replace(" ", "_")
+        if file_type_name == "RECON":
+            self.add_recon(image_stack)
+        elif file_type_name == "IMAGES":
+            self.stacks.append(image_stack)
+        else:
+            file_type = getattr(FILE_TYPES, file_type_name)
+            self.set_stack(file_type, image_stack)
+
     @property
     def is_processed(self) -> bool:
         """

mantidimaging/gui/windows/main/presenter.py

@@ -16,7 +16,7 @@
 from qt_material import apply_stylesheet
 
 from mantidimaging.core.data import ImageStack
-from mantidimaging.core.data.dataset import MixedDataset, _get_stack_data_type, Dataset
+from mantidimaging.core.data.dataset import _get_stack_data_type, Dataset
 from mantidimaging.core.io.loader.loader import create_loading_parameters_for_file_path
 from mantidimaging.core.io.utility import find_projection_closest_to_180, THRESHOLD_180
 from mantidimaging.core.utility.data_containers import ProjectionAngles
@@ -428,7 +428,7 @@ def add_180_deg_file_to_dataset(self, dataset_id: uuid.UUID, _180_deg_file: str)
             self.add_child_item_to_tree_view(dataset_id, _180_deg.id, "180")
         else:
             self.replace_child_item_id(dataset_id, existing_180_id, _180_deg.id)
-            self._delete_stack(existing_180_id)
+            self._delete_stack_visualiser(existing_180_id)
 
         self.view.model_changed.emit()
 
@@ -576,7 +576,7 @@ def _delete_container(self, container_id: uuid.UUID) -> None:
         removed_stack_ids = self.model.remove_container(container_id)
         for stack_id in removed_stack_ids:
             if stack_id in self.stack_visualisers:
-                self._delete_stack(stack_id)
+                self._delete_stack_visualiser(stack_id)
 
         # If the container_id provided is not a stack id then we remove the entire container from the tree view,
         # otherwise we remove the individual stacks that were deleted
@@ -586,7 +586,7 @@ def _delete_container(self, container_id: uuid.UUID) -> None:
 
         self.view.model_changed.emit()
 
-    def _delete_stack(self, stack_id: uuid.UUID) -> None:
+    def _delete_stack_visualiser(self, stack_id: uuid.UUID) -> None:
         """
         Deletes a stack and frees memory.
         :param stack_id: The ID of the stack to delete.
@@ -631,7 +631,7 @@ def add_sinograms_to_dataset_and_update_view(self, sino_stack: ImageStack, origi
         parent_id = self.model.get_parent_dataset(original_stack_id)
         prev_sino = self.model.datasets[parent_id].sinograms
         if prev_sino is not None:
-            self._delete_stack(prev_sino.id)
+            self._delete_stack_visualiser(prev_sino.id)
         self.model.datasets[parent_id].sinograms = sino_stack
         self._add_sinograms_to_tree_view(sino_stack.id, parent_id)
         self.create_single_tabbed_images_stack(sino_stack)
@@ -684,36 +684,15 @@ def _add_images_to_existing_dataset(self) -> None:
         assert dataset is not None
 
         new_images = self.view.add_to_dataset_dialog.presenter.images
+        images_type = self.view.add_to_dataset_dialog.images_type
 
-        if self.view.add_to_dataset_dialog.images_type == RECON_TEXT:
-            self._add_recon_to_dataset_and_tree_view(dataset, new_images)
-        elif isinstance(dataset, MixedDataset):
-            self._add_images_to_existing_mixed_dataset(dataset, new_images)
-        else:
-            self._add_images_to_existing_strict_dataset(dataset, new_images,
-                                                        self.view.add_to_dataset_dialog.images_type)
+        dataset.set_stack_by_type_name(images_type, new_images)
 
         self.create_single_tabbed_images_stack(new_images)
+        self.update_dataset_tree()
+        self._close_unused_visualisers()
         self.view.model_changed.emit()
 
-    def _add_recon_to_dataset_and_tree_view(self, dataset: Dataset, recon: ImageStack) -> None:
-        """
-        Adds a recon to the dataset and updates the tree view.
-        :param dataset: The dataset.
-        :param recon: The recon ImageStack.
-        """
-        dataset.add_recon(recon)
-        self.add_recon_item_to_tree_view(dataset.id, recon.id, recon.name)
-
-    def _add_images_to_existing_mixed_dataset(self, dataset: Dataset, new_images: ImageStack) -> None:
-        """
-        Updates the stack list in the mixed dataset and updates the tree view.
-        :param dataset: The MixedDataset to update.
-        :param new_images: The new images to add.
-        """
-        dataset.add_stack(new_images)
-        self.add_child_item_to_tree_view(dataset.id, new_images.id, new_images.name)
-
     def _add_images_to_existing_strict_dataset(self, dataset: Dataset, new_images: ImageStack, stack_type: str) -> None:
         """
         Adds or replaces images in a StrictDataset and updates the tree view if required.
@@ -722,21 +701,15 @@ def _add_images_to_existing_strict_dataset(self, dataset: Dataset, new_images: I
         """
         image_attr = stack_type.replace(" ", "_").lower()
         new_images.name = self._create_strict_dataset_stack_name(stack_type, dataset.name)
-
-        if getattr(dataset, image_attr) is None:
-            # the image type doesn't exist in the dataset
-            self.add_child_item_to_tree_view(dataset.id, new_images.id, stack_type)
-
-        else:
-            # the image type already exists in the dataset and needs to be replaced
-            prev_images_id = getattr(dataset, image_attr).id
-            if image_attr == "sample" and dataset.proj180deg:
-                self._delete_stack(dataset.proj180deg.id)
-                self.remove_item_from_tree_view(dataset.proj180deg.id)
-            self.replace_child_item_id(dataset.id, prev_images_id, new_images.id)
-            self._delete_stack(prev_images_id)
-
         setattr(dataset, image_attr, new_images)
+        self.update_dataset_tree()
+
+    def _close_unused_visualisers(self):
+        visualisers = set(self.stack_visualisers.keys())
+        stacks = {stack.id for stack in self.get_all_stacks()}
+        removed = visualisers - stacks
+        for stack_id in removed:
+            self._delete_stack_visualiser(stack_id)
 
     def _move_stack(self, origin_dataset_id: uuid.UUID, stack_id: uuid.UUID, destination_stack_type: str,
                     destination_dataset_id: uuid.UUID) -> None:
@@ -757,20 +730,9 @@ def _move_stack(self, origin_dataset_id: uuid.UUID, stack_id: uuid.UUID, destina
                 f"Unable to find destination dataset with ID {destination_dataset_id} when attempting to move stack")
 
         stack_to_move = self.get_stack(stack_id)
-        self.remove_item_from_tree_view(stack_id)
-
-        if destination_stack_type is RECON_TEXT:
-            self._add_recon_to_dataset_and_tree_view(destination_dataset, stack_to_move)
-        elif isinstance(destination_dataset, MixedDataset):
-            self._add_images_to_existing_mixed_dataset(destination_dataset, stack_to_move)
-        else:
-            assert self.view.move_stack_dialog is not None
-            data_type = self.view.move_stack_dialog.destination_stack_type
-            self._add_images_to_existing_strict_dataset(destination_dataset, stack_to_move, data_type)
-            stack_to_move.name = self._create_strict_dataset_stack_name(data_type, destination_dataset.name)
-
         origin_dataset.delete_stack(stack_id)
-        self.view.model_changed.emit()
+        self._add_images_to_existing_strict_dataset(destination_dataset, stack_to_move, destination_stack_type)
+        self._close_unused_visualisers()
 
     @staticmethod
     def _create_strict_dataset_stack_name(stack_type: str, dataset_name: str) -> str:

mantidimaging/gui/windows/main/test/presenter_test.py

@@ -9,6 +9,7 @@
 from unittest.mock import patch, call
 
 import numpy as np
+from parameterized import parameterized
 
 from mantidimaging.core.data import ImageStack
 from mantidimaging.core.data.dataset import StrictDataset, MixedDataset, Dataset
@@ -621,11 +622,11 @@ def test_add_sinograms_to_dataset_with_no_sinograms_and_update_view(self):
         dataset_item_mock.id = ds.id
         self.view.get_sinograms_item.return_value = None
         self.presenter.create_single_tabbed_images_stack = mock.Mock()
-        self.presenter._delete_stack = mock.Mock()
+        self.presenter._delete_stack_visualiser = mock.Mock()
 
         self.presenter.add_sinograms_to_dataset_and_update_view(sinograms, ds.sample.id)
         self.model.get_parent_dataset.assert_called_once_with(ds.sample.id)
-        self.presenter._delete_stack.assert_not_called()
+        self.presenter._delete_stack_visualiser.assert_not_called()
         self.assertIs(ds.sinograms, sinograms)
         self.view.get_dataset_tree_view_item.assert_called_once_with(ds.id)
         self.view.get_sinograms_item.assert_called_once_with(dataset_item_mock)
@@ -644,10 +645,10 @@ def test_add_sinograms_to_dataset_with_existing_sinograms_and_update_view(self):
         dataset_item_mock.id = ds.id
         sinograms_item_mock = self.view.get_sinograms_item.return_value
         self.presenter.create_single_tabbed_images_stack = mock.Mock()
-        self.presenter._delete_stack = mock.Mock()
+        self.presenter._delete_stack_visualiser = mock.Mock()
 
         self.presenter.add_sinograms_to_dataset_and_update_view(new_sinograms, ds.sample.id)
-        self.presenter._delete_stack.assert_called_once_with(existing_sinograms.id)
+        self.presenter._delete_stack_visualiser.assert_called_once_with(existing_sinograms.id)
         self.assertIs(ds.sinograms, new_sinograms)
         self.view.get_dataset_tree_view_item.assert_called_once_with(ds.id)
         self.view.get_sinograms_item.assert_called_once_with(dataset_item_mock)
@@ -764,87 +765,36 @@ def test_show_add_stack_to_dataset_dialog_called_with_stack_id(self):
         self.presenter.notify(Notification.SHOW_ADD_STACK_DIALOG, container_id="stack-id")
         self.view.show_add_stack_to_existing_dataset_dialog.assert_called_once_with(dataset_id)
 
-    def test_add_new_stack_to_strict_dataset(self):
+    @parameterized.expand(["Sample", "Flat Before", "Flat After", "Dark Before", "Dark After", "Recon", "Images"])
+    def test_add_new_stack_to_dataset(self, images_type):
         self.dataset.flat_before = None
         self.model.datasets[self.dataset.id] = self.dataset
 
         self.view.add_to_dataset_dialog = mock.Mock()
         self.view.add_to_dataset_dialog.dataset_id = self.dataset.id
         self.view.add_to_dataset_dialog.presenter.images = new_images = generate_images()
-        self.view.add_to_dataset_dialog.images_type = images_type = "Flat Before"
+        self.view.add_to_dataset_dialog.images_type = images_type
 
-        self.presenter.add_child_item_to_tree_view = mock.Mock()
         self.presenter.create_single_tabbed_images_stack = mock.Mock()
+        self.presenter._close_unused_visualisers = mock.Mock()
+        self.view.add_toplevel_item_to_dataset_tree_widget.return_value = mock_top_item = mock.Mock()
+        self.view.add_item_to_dataset_tree_widget.return_value = mock_top_item
 
-        self.presenter._add_images_to_existing_dataset()
-        self.presenter.add_child_item_to_tree_view.assert_called_once_with(self.dataset.id, new_images.id, images_type)
-        self.presenter.create_single_tabbed_images_stack.assert_called_once_with(new_images)
-        self.view.model_changed.emit.assert_called_once()
-        self.assertIs(self.dataset.flat_before, new_images)
-
-    def test_replace_existing_stack_in_strict_dataset(self):
-        self.model.datasets[self.dataset.id] = self.dataset
-        prev_images_id = self.dataset.flat_before.id
+        treeview_label = images_type
+        if images_type == "Sample":
+            treeview_label = "Projections"
+        elif images_type in ["Recon", "Images"]:
+            treeview_label = new_images.name
 
-        self.view.add_to_dataset_dialog = mock.Mock()
-        self.view.add_to_dataset_dialog.dataset_id = self.dataset.id
-        self.view.add_to_dataset_dialog.presenter.images = new_images = generate_images()
-        self.view.add_to_dataset_dialog.images_type = "Flat Before"
+        self.presenter._add_images_to_existing_dataset()
 
-        self.presenter.create_single_tabbed_images_stack = mock.Mock()
-        self.presenter.replace_child_item_id = mock.Mock()
-        self.presenter._delete_stack = mock.Mock()
+        self.assertIn(call(treeview_label, new_images.id, mock_top_item),
+                      self.view.add_item_to_dataset_tree_widget.mock_calls)
 
-        self.presenter.notify(Notification.DATASET_ADD)
-        self.presenter.replace_child_item_id.assert_called_once_with(self.dataset.id, prev_images_id, new_images.id)
-        self.presenter._delete_stack.assert_called_once_with(prev_images_id)
         self.presenter.create_single_tabbed_images_stack.assert_called_once_with(new_images)
         self.view.model_changed.emit.assert_called_once()
-        self.assertIs(self.dataset.flat_before, new_images)
-
-    def test_add_stack_to_mixed_dataset(self):
-        mixed_dataset = MixedDataset()
-        self.model.datasets[mixed_dataset.id] = mixed_dataset
-
-        self.view.add_to_dataset_dialog = mock.Mock()
-        self.view.add_to_dataset_dialog.dataset_id = mixed_dataset.id
-        self.view.add_to_dataset_dialog.presenter.images = new_images = generate_images()
-        self.presenter.add_child_item_to_tree_view = mock.Mock()
-
-        self.presenter._add_images_to_existing_dataset()
-        self.assertIn(new_images, mixed_dataset.all)
-        self.presenter.add_child_item_to_tree_view.assert_called_once_with(mixed_dataset.id, new_images.id,
-                                                                           new_images.name)
-
-    def test_add_recon_to_mixed_dataset(self):
-        mixed_dataset = MixedDataset(stacks=[generate_images()])
-        recon = generate_images()
-        recon.name = "recon-name"
-
-        self.view.add_to_dataset_dialog.images_type = RECON_TEXT
-        self.presenter.add_recon_item_to_tree_view = mock.Mock()
-        self.presenter.get_dataset = mock.Mock(return_value=mixed_dataset)
-        self.view.add_to_dataset_dialog.presenter.images = recon
-
-        self.presenter._add_images_to_existing_dataset()
-
-        self.assertIn(recon, mixed_dataset.recons.stacks)
-        self.presenter.add_recon_item_to_tree_view.assert_called_once_with(mixed_dataset.id, recon.id, recon.name)
-
-    def test_add_recon_to_strict_dataset(self):
-        self.presenter.add_recon_item_to_tree_view = mock.Mock()
-        self.view.add_to_dataset_dialog.images_type = RECON_TEXT
-
-        recon = generate_images()
-        recon.name = "recon-name"
-        self.presenter.add_recon_item_to_tree_view = mock.Mock()
-        self.presenter.get_dataset = mock.Mock(return_value=self.dataset)
-        self.view.add_to_dataset_dialog.presenter.images = recon
-
-        self.presenter._add_images_to_existing_dataset()
-
-        self.assertIn(recon, self.dataset.recons.stacks)
-        self.presenter.add_recon_item_to_tree_view.assert_called_once_with(self.dataset.id, recon.id, recon.name)
+        self.assertIn(new_images, self.dataset.all)
+        self.presenter._close_unused_visualisers.assert_called_once()
 
     def test_select_tree_widget_item(self):
         tree_widget_item = mock.Mock()
@@ -949,7 +899,7 @@ def test_stack_moved_to_recon(self):
 
         self.presenter.get_stack = mock.Mock(return_value=stack_to_move)
         self.presenter.remove_item_from_tree_view = mock.Mock()
-        self.presenter._add_recon_to_dataset_and_tree_view = mock.Mock()
+        self.presenter._add_images_to_existing_strict_dataset = mock.Mock()
         self.presenter.get_dataset = mock.Mock(side_effect=[origin_dataset, destination_dataset])
 
         self.presenter.notify(Notification.MOVE_STACK,
@@ -958,7 +908,8 @@ def test_stack_moved_to_recon(self):
                               destination_stack_type=RECON_TEXT,
                               destination_dataset_id=destination_dataset_id)
         self.presenter.get_stack.assert_called_once_with(stack_to_move.id)
-        self.presenter._add_recon_to_dataset_and_tree_view.assert_called_once_with(destination_dataset, stack_to_move)
+        self.presenter._add_images_to_existing_strict_dataset.assert_called_once_with(
+            destination_dataset, stack_to_move, RECON_TEXT)
 
         self.assertNotIn(stack_to_move, origin_dataset)
 
@@ -971,34 +922,34 @@ def test_stack_moved_to_mixed_dataset_images(self):
         self.presenter.get_dataset = mock.Mock(side_effect=[origin_dataset, destination_dataset])
         self.presenter.get_stack = mock.Mock(return_value=stack_to_move)
         self.presenter.remove_item_from_tree_view = mock.Mock()
-        self.presenter._add_images_to_existing_mixed_dataset = mock.Mock()
+        self.presenter._add_images_to_existing_strict_dataset = mock.Mock()
 
         self.presenter._move_stack(origin_dataset.id, stack_to_move.id, "Images", destination_dataset_id)
         self.presenter.get_stack.assert_called_once_with(stack_to_move.id)
-        self.presenter._add_images_to_existing_mixed_dataset.assert_called_once_with(destination_dataset, stack_to_move)
+        self.presenter._add_images_to_existing_strict_dataset.assert_called_once_with(
+            destination_dataset, stack_to_move, "Images")
 
         self.assertNotIn(stack_to_move, origin_dataset)
 
     def test_move_stack_to_strict_dataset(self):
         stack_to_move = generate_images()
+        print(f"{stack_to_move.name=}")
         origin_dataset = MixedDataset(stacks=[stack_to_move])
         destination_dataset = StrictDataset(sample=generate_images())
         self.presenter.get_dataset = mock.Mock(side_effect=[origin_dataset, destination_dataset])
         self.presenter.get_stack = mock.Mock(return_value=stack_to_move)
 
         self.presenter.remove_item_from_tree_view = mock.Mock()
-        self.presenter._add_images_to_existing_strict_dataset = mock.Mock()
+        self.presenter.update_dataset_tree = mock.Mock()
 
         self.view.move_stack_dialog = mock.Mock()
         self.view.move_stack_dialog.destination_stack_type = data_type = "Flat After"
-        self._add_images_to_existing_strict_dataset = mock.Mock()
         new_stack_name = "New Dataset Flat After"
         self.presenter._create_strict_dataset_stack_name = mock.Mock(return_value=new_stack_name)
 
         self.presenter._move_stack(origin_dataset.id, stack_to_move.id, data_type, destination_dataset.id)
         self.presenter.get_stack.assert_called_once_with(stack_to_move.id)
-        self.presenter._add_images_to_existing_strict_dataset.assert_called_once_with(
-            destination_dataset, stack_to_move, data_type)
+        self.presenter.update_dataset_tree.assert_called_once()
 
         self.assertNotIn(stack_to_move, origin_dataset)
         assert stack_to_move.name == new_stack_name
@@ -1107,6 +1058,21 @@ def test_update_dataset_tree_datasets_with_image_stacks(self):
         ]
         self.assertListEqual(expected_calls, self.view.add_item_to_dataset_tree_widget.mock_calls)
 
+    def test_close_unused_visualisers(self):
+        stacks_ids = [uuid.uuid4() for _ in range(5)]
+        visualisers = {id: mock.Mock() for id in stacks_ids}
+        stacks = [mock.Mock(id=id) for id in stacks_ids]
+        self.presenter.stack_visualisers = visualisers
+        self.presenter.get_all_stacks = mock.Mock(return_value=stacks)
+        self.presenter._delete_stack_visualiser = mock.Mock()
+
+        self.presenter._close_unused_visualisers()
+        self.presenter._delete_stack_visualiser.assert_not_called()
+
+        stacks.pop(2)
+        self.presenter._close_unused_visualisers()
+        self.presenter._delete_stack_visualiser.assert_called_with(stacks_ids[2])
+
 
 if __name__ == '__main__':
     unittest.main()