Unify datasets part 5 (#2323)

Author: MikeSullivan7

mantidimaging/core/data/dataset.py

@@ -22,7 +22,7 @@ def remove_nones(image_stacks: list[ImageStack | None]) -> list[ImageStack]:
     return [image_stack for image_stack in image_stacks if image_stack is not None]
 
 
-class BaseDataset:
+class Dataset:
 
     def __init__(
         self,
@@ -185,15 +185,15 @@ def is_processed(self) -> bool:
         return False
 
 
-class MixedDataset(BaseDataset):
+class MixedDataset(Dataset):
     pass
 
 
-class StrictDataset(BaseDataset):
+class StrictDataset(Dataset):
     pass
 
 
-def _get_stack_data_type(stack_id: uuid.UUID, dataset: BaseDataset) -> str:
+def _get_stack_data_type(stack_id: uuid.UUID, dataset: Dataset) -> str:
     """
     Find the data type as a string of a stack.
     :param stack_id: The ID of the stack.

mantidimaging/core/data/test/dataset_test.py

@@ -9,7 +9,7 @@
 import numpy as np
 
 from mantidimaging.core.data import ImageStack
-from mantidimaging.core.data.dataset import BaseDataset, _get_stack_data_type
+from mantidimaging.core.data.dataset import Dataset, _get_stack_data_type
 from mantidimaging.core.utility.data_containers import ProjectionAngles, FILE_TYPES
 from mantidimaging.test_helpers.unit_test_helper import generate_images
 
@@ -19,11 +19,11 @@ def _make_standard_dataset(shape=(2, 5, 5)):
     image_stacks = [generate_images(shape) for _ in range(6)]
     image_stacks[0].name = "samplename"
 
-    ds = BaseDataset(sample=image_stacks[0],
-                     flat_before=image_stacks[1],
-                     flat_after=image_stacks[2],
-                     dark_before=image_stacks[3],
-                     dark_after=image_stacks[4])
+    ds = Dataset(sample=image_stacks[0],
+                 flat_before=image_stacks[1],
+                 flat_after=image_stacks[2],
+                 dark_before=image_stacks[3],
+                 dark_after=image_stacks[4])
     ds.proj180deg = image_stacks[5]
     return ds, image_stacks
 
@@ -41,19 +41,19 @@ def _set_fake_projection_angles(image_stack: ImageStack):
 class DatasetTest(unittest.TestCase):
 
     def test_create_dataset(self):
-        ds = BaseDataset()
+        ds = Dataset()
         self.assertIsInstance(ds.id, uuid.UUID)
         self.assertEqual(ds.name, "")
 
     def test_prevent_positional_parameters(self):
-        self.assertRaises(TypeError, BaseDataset, mock.Mock())
+        self.assertRaises(TypeError, Dataset, mock.Mock())
 
     def test_dataset_name(self):
-        ds = BaseDataset(name="a_dataset")
+        ds = Dataset(name="a_dataset")
         self.assertEqual(ds.name, "a_dataset")
 
     def test_add_recon(self):
-        ds = BaseDataset()
+        ds = Dataset()
         recons = [generate_images() for _ in range(3)]
         [ds.add_recon(r) for r in recons]
 
@@ -62,7 +62,7 @@ def test_add_recon(self):
         self.assertEqual(3, len(ds.recons))
 
     def test_delete_recon(self):
-        ds = BaseDataset()
+        ds = Dataset()
         recons = [generate_images() for _ in range(3)]
         [ds.add_recon(r) for r in recons]
 
@@ -71,36 +71,36 @@ def test_delete_recon(self):
         self.assertNotIn(recons[-1], ds.all)
 
     def test_delete_failure(self):
-        ds = BaseDataset()
+        ds = Dataset()
         with self.assertRaises(KeyError):
             ds.delete_stack("nonexistent-id")
 
     def test_delete_all_recons(self):
-        ds = BaseDataset()
+        ds = Dataset()
         recons = [generate_images() for _ in range(3)]
         [ds.add_recon(r) for r in recons]
         ds.delete_recons()
         self.assertListEqual(ds.recons.stacks, [])
 
     def test_sinograms(self):
-        ds = BaseDataset()
+        ds = Dataset()
         ds.sinograms = sinograms = generate_images()
         self.assertIs(ds.sinograms, sinograms)
 
     def test_delete_sinograms(self):
-        ds = BaseDataset()
+        ds = Dataset()
         ds.sinograms = sinograms = generate_images()
         ds.delete_stack(sinograms.id)
         self.assertIsNone(ds.sinograms)
 
     def test_stacks_in_all(self):
         image_stacks = [mock.Mock() for _ in range(3)]
-        ds = BaseDataset(stacks=image_stacks)
+        ds = Dataset(stacks=image_stacks)
         self.assertListEqual(ds.all, image_stacks)
 
     def test_sample_in_all(self):
         image_sample = mock.Mock(proj180deg=None)
-        ds = BaseDataset(sample=image_sample)
+        ds = Dataset(sample=image_sample)
         self.assertCountEqual(ds.all, [image_sample])
 
     def test_all_for_full_dataset(self):
@@ -111,27 +111,27 @@ def test_all_for_full_dataset(self):
 
     def test_delete_stack_from_stacks_list(self):
         image_stacks = [mock.Mock() for _ in range(3)]
-        ds = BaseDataset(stacks=image_stacks)
+        ds = Dataset(stacks=image_stacks)
         prev_stacks = image_stacks.copy()
         ds.delete_stack(image_stacks[-1].id)
         self.assertListEqual(ds.all, prev_stacks[:-1])
 
     def test_get_stack_data_type_returns_recon(self):
         recon = generate_images()
         recon_id = recon.id
-        dataset = BaseDataset()
+        dataset = Dataset()
         dataset.recons.append(recon)
         self.assertEqual(_get_stack_data_type(recon_id, dataset), "Recon")
 
     def test_get_stack_data_type_returns_images(self):
         images = generate_images()
         images_id = images.id
-        dataset = BaseDataset(stacks=[images])
+        dataset = Dataset(stacks=[images])
         self.assertEqual(_get_stack_data_type(images_id, dataset), "Images")
 
     def test_attribute_not_set_returns_none(self):
         sample = mock.Mock()
-        dataset = BaseDataset(sample=sample)
+        dataset = Dataset(sample=sample)
 
         self.assertIsNone(dataset.flat_before)
         self.assertIsNone(dataset.flat_after)
@@ -140,7 +140,7 @@ def test_attribute_not_set_returns_none(self):
 
     def test_set_flat_before(self):
         sample = mock.Mock()
-        dataset = BaseDataset(sample=sample)
+        dataset = Dataset(sample=sample)
         flat_before = mock.Mock(id="1234")
         dataset.flat_before = flat_before
         self.assertIs(flat_before, dataset.flat_before)
@@ -258,21 +258,21 @@ def test_delete_dark_after(self):
         self.assertIsNone(ds.dark_after)
 
     def test_set_stack_by_type_sample(self):
-        ds = BaseDataset()
+        ds = Dataset()
         sample = mock.Mock()
         ds.set_stack(FILE_TYPES.SAMPLE, sample)
 
         self.assertEqual(ds.sample, sample)
 
     def test_set_stack_by_type_flat_before(self):
-        ds = BaseDataset()
+        ds = Dataset()
         stack = mock.Mock()
         ds.set_stack(FILE_TYPES.FLAT_BEFORE, stack)
 
         self.assertEqual(ds.flat_before, stack)
 
     def test_set_stack_by_type_180(self):
-        ds = BaseDataset()
+        ds = Dataset()
         sample = mock.Mock()
         stack = mock.Mock()
         ds.set_stack(FILE_TYPES.SAMPLE, sample)
@@ -281,10 +281,10 @@ def test_set_stack_by_type_180(self):
         self.assertEqual(ds.proj180deg, stack)
 
     def test_processed_is_true(self):
-        ds = BaseDataset(sample=generate_images())
+        ds = Dataset(sample=generate_images())
         ds.sample.record_operation("", "")
         self.assertTrue(ds.is_processed)
 
     def test_processed_is_false(self):
-        ds = BaseDataset(sample=generate_images())
+        ds = Dataset(sample=generate_images())
         self.assertFalse(ds.is_processed)

mantidimaging/eyes_tests/main_window_test.py

@@ -87,11 +87,6 @@ def test_show_add_stack_to_existing_dataset_dialog_for_strict_dataset(self):
 
         self.check_target(widget=self.imaging.add_to_dataset_dialog)
 
-    def test_show_add_stack_to_existing_dataset_dialog_for_mixed_dataset(self):
-        self._create_mixed_dataset()
-        self.imaging.show_add_stack_to_existing_dataset_dialog(list(self.imaging.presenter.all_dataset_ids)[0])
-        self.check_target(widget=self.imaging.add_to_dataset_dialog)
-
     def test_clicking_tab_changes_tree_view_selection(self):
         self._load_strict_data_set()
         sample_vis = self._load_strict_data_set()

mantidimaging/gui/windows/add_images_to_dataset_dialog/test/view_test.py

@@ -22,7 +22,7 @@ def setUp(self) -> None:
             self.main_window = MainWindowView()
         self.dataset_id = uuid.uuid4()
         self.dataset_name = "dataset-name"
-        self.view = AddImagesToDatasetDialog(self.main_window, self.dataset_id, True, self.dataset_name)
+        self.view = AddImagesToDatasetDialog(self.main_window, self.dataset_id, self.dataset_name)
         self.view.presenter = self.presenter = mock.MagicMock()
 
     def test_on_accepted(self):

mantidimaging/gui/windows/add_images_to_dataset_dialog/view.py

@@ -15,18 +15,15 @@ class AddImagesToDatasetDialog(BaseDialogView):
     filePathLineEdit: QLineEdit
     datasetNameText: QLabel
 
-    def __init__(self, parent, dataset_id: uuid.UUID, strict_dataset: bool, dataset_name: str):
+    def __init__(self, parent, dataset_id: uuid.UUID, dataset_name: str):
         super().__init__(parent, 'gui/ui/add_to_dataset.ui')
 
         self.parent_view = parent
         self.presenter = AddImagesToDatasetPresenter(self)
         self._dataset_id = dataset_id
 
-        if strict_dataset:
-            self.imageTypeComboBox.addItems(
-                ["Sample", "Flat Before", "Flat After", "Dark Before", "Dark After", "Recon"])
-        else:
-            self.imageTypeComboBox.addItems(["Images", "Recon"])
+        self.imageTypeComboBox.addItems(
+            ["Sample", "Flat Before", "Flat After", "Dark Before", "Dark After", "Recon", "Images"])
 
         self.datasetNameText.setText(dataset_name)
         self.chooseFileButton.clicked.connect(self.choose_file_path)

mantidimaging/gui/windows/main/model.py

@@ -7,7 +7,7 @@
 from typing import NoReturn, TYPE_CHECKING
 
 from mantidimaging.core.data import ImageStack
-from mantidimaging.core.data.dataset import StrictDataset, MixedDataset
+from mantidimaging.core.data.dataset import StrictDataset, MixedDataset, Dataset
 from mantidimaging.core.io import loader, saver
 from mantidimaging.core.io.filenames import FilenameGroup
 from mantidimaging.core.io.loader.loader import LoadingParameters, ImageParameters
@@ -27,7 +27,7 @@ class MainWindowModel:
 
     def __init__(self) -> None:
         super().__init__()
-        self.datasets: dict[uuid.UUID, MixedDataset | StrictDataset] = {}
+        self.datasets: dict[uuid.UUID, Dataset] = {}
 
     def get_images_by_uuid(self, images_uuid: uuid.UUID) -> ImageStack | None:
         for dataset in self.datasets.values():
@@ -102,8 +102,10 @@ def get_existing_180_id(self, dataset_id: uuid.UUID) -> uuid.UUID | None:
         :return: The 180 ID if found, None otherwise.
         """
         dataset = self.datasets.get(dataset_id)
-        if not isinstance(dataset, StrictDataset):
-            raise RuntimeError(f"Failed to get StrictDataset with ID {dataset_id}")
+        if not dataset:
+            raise RuntimeError(f"Failed to get Dataset with ID {dataset_id}")
+        if not dataset.sample:
+            raise RuntimeError(f"Dataset with ID {dataset_id} does not have a sample")
 
         if isinstance(dataset.proj180deg, ImageStack):
             return dataset.proj180deg.id
@@ -116,13 +118,11 @@ def add_180_deg_to_dataset(self, dataset_id: uuid.UUID, _180_deg_file: str) -> I
         :param _180_deg_file: The location of the 180 projection.
         :return: The loaded 180 ImageStack object.
         """
-        if dataset_id in self.datasets:
-            dataset = self.datasets[dataset_id]
-        else:
+        dataset = self.datasets.get(dataset_id)
+        if not dataset:
             raise RuntimeError(f"Failed to get Dataset with ID {dataset_id}")
-
-        if not isinstance(dataset, StrictDataset):
-            raise RuntimeError(f"Wrong dataset type passed to add 180 method: {dataset_id}")
+        if not dataset.sample:
+            raise RuntimeError(f"Dataset with ID {dataset_id} does not have a sample")
 
         _180_deg = loader.load_stack_from_group(FilenameGroup.from_file(_180_deg_file))
         dataset.proj180deg = _180_deg
@@ -201,7 +201,7 @@ def remove_container(self, container_id: uuid.UUID) -> list[uuid.UUID]:
                     return ids_to_remove
         self.raise_error_when_images_not_found(container_id)
 
-    def add_dataset_to_model(self, dataset: StrictDataset | MixedDataset) -> None:
+    def add_dataset_to_model(self, dataset: Dataset) -> None:
         self.datasets[dataset.id] = dataset
 
     @property

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 StrictDataset, MixedDataset, _get_stack_data_type
+from mantidimaging.core.data.dataset import StrictDataset, MixedDataset, _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
@@ -416,7 +416,7 @@ def stack_visualiser_list(self) -> list[StackId]:
         return sorted(stacks, key=lambda x: x.name)
 
     @property
-    def datasets(self) -> Iterable[MixedDataset | StrictDataset]:
+    def datasets(self) -> Iterable[Dataset]:
         return self.model.datasets.values()
 
     @property
@@ -442,7 +442,7 @@ def all_stack_ids(self) -> Iterable[uuid.UUID]:
     def stack_visualiser_names(self) -> list[str]:
         return [widget.windowTitle() for widget in self.stack_visualisers.values()]
 
-    def get_dataset(self, dataset_id: uuid.UUID) -> MixedDataset | StrictDataset | None:
+    def get_dataset(self, dataset_id: uuid.UUID) -> Dataset | None:
         return self.model.datasets.get(dataset_id)
 
     def get_stack_visualiser(self, stack_id: uuid.UUID) -> StackVisualiserView:
@@ -757,7 +757,7 @@ def _add_images_to_existing_dataset(self) -> None:
         self.create_single_tabbed_images_stack(new_images)
         self.view.model_changed.emit()
 
-    def _add_recon_to_dataset_and_tree_view(self, dataset: MixedDataset | StrictDataset, recon: ImageStack) -> None:
+    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.
@@ -766,7 +766,7 @@ def _add_recon_to_dataset_and_tree_view(self, dataset: MixedDataset | StrictData
         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: MixedDataset, new_images: ImageStack) -> None:
+    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.
@@ -775,8 +775,7 @@ def _add_images_to_existing_mixed_dataset(self, dataset: MixedDataset, new_image
         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: StrictDataset, new_images: ImageStack,
-                                               stack_type: str) -> None:
+    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.
         :param dataset: The StrictDataset to change.

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

@@ -10,7 +10,7 @@
 import numpy as np
 
 from mantidimaging.core.data import ImageStack
-from mantidimaging.core.data.dataset import StrictDataset, MixedDataset
+from mantidimaging.core.data.dataset import StrictDataset, MixedDataset, Dataset
 from mantidimaging.core.io.loader.loader import LoadingParameters, ImageParameters
 from mantidimaging.core.utility.data_containers import ProjectionAngles, FILE_TYPES, Indices
 from mantidimaging.gui.windows.main import MainWindowModel
@@ -487,21 +487,21 @@ def test_no_dataset_with_180_raises(self):
             self.model.get_existing_180_id("bad-id")
 
     def test_wrong_dataset_type_for_180_raises(self):
-        md = MixedDataset()
+        md = Dataset()
         self.model.add_dataset_to_model(md)
 
         with self.assertRaises(RuntimeError):
             self.model.get_existing_180_id(md.id)
 
     def test_get_existing_180_id_finds_id(self):
-        sd = StrictDataset(sample=generate_images((5, 20, 20)))
+        sd = Dataset(sample=generate_images((5, 20, 20)))
         sd.proj180deg = _180 = generate_images((1, 20, 20))
         self.model.add_dataset_to_model(sd)
 
         assert self.model.get_existing_180_id(sd.id) == _180.id
 
     def test_get_existing_id_returns_none_for_dataset_without_180(self):
-        sd = StrictDataset(sample=generate_images((5, 20, 20)))
+        sd = Dataset(sample=generate_images((5, 20, 20)))
         self.model.add_dataset_to_model(sd)
 
         self.assertIsNone(self.model.get_existing_180_id(sd.id))