Merge pull request #255 from bessagroup/mpvanderschelling/issue250

Mpvanderschelling/issue250

src/f3dasm/_src/experimentdata/_jobqueue.py

@@ -167,6 +167,24 @@ def reset(self) -> None:
         """Resets the job queue."""
         self.jobs = pd.Series(dtype='string')
 
+    #                                                                    Select
+    # =========================================================================
+
+    def select_all(self, status: str) -> _JobQueue:
+        """Selects all jobs with a certain status.
+
+        Parameters
+        ----------
+        status : str
+            Status of the jobs to select
+
+        Returns
+        -------
+        JobQueue
+            JobQueue object containing the selected jobs.
+        """
+        return _JobQueue(self.jobs[self.jobs == status])
+
     #                                                                    Export
     # =========================================================================
 

src/f3dasm/_src/experimentdata/experimentdata.py

@@ -377,6 +377,34 @@ def select(self, indices: int | Iterable[int]) -> ExperimentData:
                               jobs=self._jobs[indices],
                               domain=self.domain, project_dir=self.project_dir)
 
+    def select_with_status(self, status: Literal['open', 'in progress',
+                                                 'finished', 'error']
+                           ) -> ExperimentData:
+        """Select a subset of the ExperimentData object with a given status
+
+        Parameters
+        ----------
+        status : Literal['open', 'in progress', 'finished', 'error']
+            The status to select.
+
+        Returns
+        -------
+        ExperimentData
+            The selected ExperimentData object with only the selected status.
+
+        Raises
+        ------
+        ValueError
+            Raised when invalid status is specified
+        """
+        if status not in [s.value for s in Status]:
+            raise ValueError(f"Invalid status {status} given. "
+                             f"\nChoose from values: "
+                             f"{', '.join([s.value for s in Status])}")
+
+        _indices = self._jobs.select_all(status).indices
+        return self.select(_indices)
+
     def get_input_data(self,
                        parameter_names: Optional[str | Iterable[str]] = None
                        ) -> ExperimentData:

tests/experimentdata/test__jobqueue.py

@@ -0,0 +1,36 @@
+import pandas as pd
+
+from f3dasm._src.experimentdata._jobqueue import _JobQueue
+
+
+def test_select_all_with_matching_status():
+    # Create a job queue with some jobs
+    job_queue = _JobQueue()
+    job_queue.jobs = pd.Series(['in progress', 'running', 'completed', 'in progress', 'failed'])
+
+    # Select all jobs with status 'in progress'
+    selected_jobs = job_queue.select_all('in progress')
+
+    # Check if the selected jobs match the expected result
+    assert (selected_jobs.jobs == ['in progress', 'in progress']).all()
+
+def test_select_all_with_no_matching_status():
+    # Create a job queue with some jobs
+    job_queue = _JobQueue()
+    job_queue.jobs = pd.Series(['in progress', 'running', 'completed', 'in progress', 'failed'])
+
+    # Select all jobs with status 'cancelled'
+    selected_jobs = job_queue.select_all('cancelled')
+
+    # Check if the selected jobs match the expected result
+    assert selected_jobs.jobs.empty
+
+def test_select_all_with_empty_job_queue():
+    # Create an empty job queue
+    job_queue = _JobQueue()
+
+    # Select all jobs with status 'in progress'
+    selected_jobs = job_queue.select_all('in progress')
+
+    # Check if the selected jobs match the expected result
+    assert selected_jobs.jobs.empty

tests/experimentdata/test_experimentdata.py

@@ -703,5 +703,29 @@ def test_iter_behaviour(experimentdata_continuous: ExperimentData):
         assert isinstance(i, ExperimentSample)
 
 
+def test_select_with_status_open(experimentdata: ExperimentData):
+    selected_data = experimentdata.select_with_status('open')
+    assert all(job == Status.OPEN for job in selected_data._jobs.jobs)
+
+
+def test_select_with_status_in_progress(experimentdata: ExperimentData):
+    selected_data = experimentdata.select_with_status('in progress')
+    assert all(job == Status.IN_PROGRESS for job in selected_data._jobs.jobs)
+
+
+def test_select_with_status_finished(experimentdata: ExperimentData):
+    selected_data = experimentdata.select_with_status('finished')
+    assert all(job == Status.FINISHED for job in selected_data._jobs.jobs)
+
+
+def test_select_with_status_error(experimentdata: ExperimentData):
+    selected_data = experimentdata.select_with_status('error')
+    assert all(job == Status.ERROR for job in selected_data._jobs.jobs)
+
+
+def test_select_with_status_invalid_status(experimentdata: ExperimentData):
+    with pytest.raises(ValueError):
+        _ = experimentdata.select_with_status('invalid_status')
+
 if __name__ == "__main__":  # pragma: no cover
-    pytest.main()
+    pytest.main()