diff --git a/src/f3dasm/_src/experimentdata/_jobqueue.py b/src/f3dasm/_src/experimentdata/_jobqueue.py index 10ecbe11..cee1e9e5 100644 --- a/src/f3dasm/_src/experimentdata/_jobqueue.py +++ b/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 # ========================================================================= diff --git a/src/f3dasm/_src/experimentdata/experimentdata.py b/src/f3dasm/_src/experimentdata/experimentdata.py index 5ae1f1c6..34855e8f 100644 --- a/src/f3dasm/_src/experimentdata/experimentdata.py +++ b/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: diff --git a/tests/experimentdata/test__jobqueue.py b/tests/experimentdata/test__jobqueue.py new file mode 100644 index 00000000..b38268bc --- /dev/null +++ b/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 diff --git a/tests/experimentdata/test_experimentdata.py b/tests/experimentdata/test_experimentdata.py index c53b0be6..37989401 100644 --- a/tests/experimentdata/test_experimentdata.py +++ b/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() \ No newline at end of file