Merge StrictDataset constructor into BaseDataset

Author: samtygier-stfc

mantidimaging/core/data/dataset.py

@@ -2,7 +2,6 @@
 # SPDX - License - Identifier: GPL-3.0-or-later
 from __future__ import annotations
 import uuid
-from dataclasses import dataclass
 
 import numpy as np
 
@@ -25,7 +24,17 @@ def remove_nones(image_stacks: list[ImageStack | None]) -> list[ImageStack]:
 
 class BaseDataset:
 
-    def __init__(self, *, name: str = "", stacks: list[ImageStack] | None = None) -> None:
+    def __init__(
+        self,
+        *,
+        name: str = "",
+        stacks: list[ImageStack] | None = None,
+        sample: ImageStack | None = None,
+        flat_before: ImageStack | None = None,
+        flat_after: ImageStack | None = None,
+        dark_before: ImageStack | None = None,
+        dark_after: ImageStack | None = None,
+    ) -> None:
         self._id: uuid.UUID = uuid.uuid4()
         self.name = name
 
@@ -34,6 +43,15 @@ def __init__(self, *, name: str = "", stacks: list[ImageStack] | None = None) ->
         stacks = [] if stacks is None else stacks
         self._stacks: list[ImageStack] = stacks
 
+        self.sample = sample
+        self.flat_before = flat_before
+        self.flat_after = flat_after
+        self.dark_before = dark_before
+        self.dark_after = dark_after
+
+        if self.name == "" and sample is not None:
+            self.name = sample.name
+
     @property
     def id(self) -> uuid.UUID:
         return self._id
@@ -89,31 +107,7 @@ class MixedDataset(BaseDataset):
     pass
 
 
-@dataclass
 class StrictDataset(BaseDataset):
-    sample: ImageStack
-    flat_before: ImageStack | None = None
-    flat_after: ImageStack | None = None
-    dark_before: ImageStack | None = None
-    dark_after: ImageStack | None = None
-
-    def __init__(self,
-                 *,
-                 sample: ImageStack,
-                 flat_before: ImageStack | None = None,
-                 flat_after: ImageStack | None = None,
-                 dark_before: ImageStack | None = None,
-                 dark_after: ImageStack | None = None,
-                 name: str = ""):
-        super().__init__(name=name)
-        self.sample = sample
-        self.flat_before = flat_before
-        self.flat_after = flat_after
-        self.dark_before = dark_before
-        self.dark_after = dark_after
-
-        if self.name == "":
-            self.name = sample.name
 
     @property
     def all(self) -> list[ImageStack]:
@@ -165,6 +159,8 @@ def proj180deg(self) -> ImageStack | None:
 
     @proj180deg.setter
     def proj180deg(self, proj180deg: ImageStack | None) -> None:
+        if self.sample is None:
+            raise RuntimeError("Can't set a 180 projection without a sample")
         self.sample.proj180deg = proj180deg
 
     def delete_stack(self, images_id: uuid.UUID) -> None:
@@ -179,6 +175,7 @@ def delete_stack(self, images_id: uuid.UUID) -> None:
         elif isinstance(self.dark_after, ImageStack) and self.dark_after.id == images_id:
             self.dark_after = None
         elif isinstance(self.proj180deg, ImageStack) and self.proj180deg.id == images_id:
+            assert self.sample is not None
             self.sample.clear_proj180deg()
         elif isinstance(self.sinograms, ImageStack) and self.sinograms.id == images_id:
             self.sinograms = None
@@ -221,7 +218,7 @@ def _get_stack_data_type(stack_id: uuid.UUID, dataset: BaseDataset) -> str:
     if stack_id in [stack.id for stack in dataset._stacks]:
         return "Images"
     if isinstance(dataset, StrictDataset):
-        if stack_id == dataset.sample.id:
+        if dataset.sample is not None and stack_id == dataset.sample.id:
             return "Sample"
         if dataset.flat_before is not None and stack_id == dataset.flat_before.id:
             return "Flat Before"

mantidimaging/core/io/saver.py

@@ -247,6 +247,7 @@ def _nexus_save(nexus_file: h5py.File, dataset: StrictDataset, sample_name: str,
     data["image_key"] = detector["image_key"]
 
     for recon in dataset.recons:
+        assert dataset.sample is not None
         assert dataset.sample.filenames is not None
         _save_recon_to_nexus(nexus_file, recon, dataset.sample.filenames[0])
 

mantidimaging/gui/windows/main/model.py

@@ -185,8 +185,10 @@ def remove_container(self, container_id: uuid.UUID) -> list[uuid.UUID]:
                     proj_180_id = None
                     # If we're deleting a sample from a StrictDataset then any linked 180 projection will also be
                     # deleted
-                    if isinstance(dataset, StrictDataset) and dataset.proj180deg and dataset.sample.id == container_id:
-                        proj_180_id = dataset.proj180deg.id
+                    if isinstance(dataset, StrictDataset) and dataset.proj180deg:
+                        assert dataset.sample is not None
+                        if dataset.sample.id == container_id:
+                            proj_180_id = dataset.proj180deg.id
                     dataset.delete_stack(container_id)
                     return [container_id, proj_180_id] if proj_180_id else [container_id]
                 if container_id == dataset.recons.id:

mantidimaging/gui/windows/main/presenter.py

@@ -241,6 +241,7 @@ def add_alternative_180_if_required(self, dataset: StrictDataset) -> None:
         Checks if the dataset has a 180 projection and tries to find an alternative if one is missing.
         :param dataset: The loaded dataset.
         """
+        assert dataset.sample is not None
         if dataset.sample.has_proj180deg() and dataset.sample.proj180deg.filenames:  # type: ignore
             return
         else:
@@ -260,6 +261,7 @@ def create_strict_dataset_stack_windows(self, dataset: StrictDataset) -> StackVi
         :param dataset: The loaded dataset.
         :return: The stack widget for the sample.
         """
+        assert dataset.sample is not None
         sample_stack_vis = self._create_lone_stack_window(dataset.sample)
         self._tabify_stack_window(sample_stack_vis)
 
@@ -357,6 +359,7 @@ def create_strict_dataset_tree_view_items(self, dataset: StrictDataset) -> None:
         Creates the tree view items for a strict dataset.
         :param dataset: The loaded dataset.
         """
+        assert dataset.sample is not None
         dataset_tree_item = self.view.create_dataset_tree_widget_item(dataset.name, dataset.id)
         self.view.create_child_tree_item(dataset_tree_item, dataset.sample.id, "Projections")