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()