diff --git a/openml/study/study.py b/openml/study/study.py index 803c6455b..57c539ec0 100644 --- a/openml/study/study.py +++ b/openml/study/study.py @@ -274,6 +274,59 @@ def __init__( # noqa: PLR0913 setups=setups, ) + def attach_runs(self, run_ids: list[int]) -> int: + """Attach runs to this study. + + Parameters + ---------- + run_ids : list[int] + List of run ids to attach to this study. + + Returns + ------- + int + The new number of linked entities in the study. + + Raises + ------ + ValueError + If the study has not been published yet. + """ + if self.id is None: + raise ValueError( + "Cannot attach runs to an unpublished study. Please publish the study first.", + ) + result = openml.study.attach_to_study(self.id, run_ids) + self.runs = (self.runs or []) + list(run_ids) + return result + + def detach_runs(self, run_ids: list[int]) -> int: + """Detach runs from this study. + + Parameters + ---------- + run_ids : list[int] + List of run ids to detach from this study. + + Returns + ------- + int + The new number of linked entities in the study. + + Raises + ------ + ValueError + If the study has not been published yet. + """ + if self.id is None: + raise ValueError( + "Cannot detach runs from an unpublished study. Please publish the study first.", + ) + result = openml.study.detach_from_study(self.id, run_ids) + if self.runs is not None: + self.runs = [run_id for run_id in self.runs if run_id not in set(run_ids)] + return result + class OpenMLBenchmarkSuite(BaseStudy): """ @@ -343,3 +396,56 @@ def __init__( # noqa: PLR0913 runs=None, setups=None, ) + + def attach_tasks(self, task_ids: list[int]) -> int: + """Attach tasks to this benchmark suite. + + Parameters + ---------- + task_ids : list[int] + List of task ids to attach to this suite. + + Returns + ------- + int + The new number of linked entities in the suite. + + Raises + ------ + ValueError + If the suite has not been published yet. + """ + if self.id is None: + raise ValueError( + "Cannot attach tasks to an unpublished suite. Please publish the suite first.", + ) + result = openml.study.attach_to_suite(self.id, task_ids) + self.tasks = (self.tasks or []) + list(task_ids) + return result + + def detach_tasks(self, task_ids: list[int]) -> int: + """Detach tasks from this benchmark suite. + + Parameters + ---------- + task_ids : list[int] + List of task ids to detach from this suite. + + Returns + ------- + int + The new number of linked entities in the suite. + + Raises + ------ + ValueError + If the suite has not been published yet. + """ + if self.id is None: + raise ValueError( + "Cannot detach tasks from an unpublished suite. Please publish the suite first.", + ) + result = openml.study.detach_from_suite(self.id, task_ids) + if self.tasks is not None: + self.tasks = [task_id for task_id in self.tasks if task_id not in set(task_ids)] + return result diff --git a/tests/test_study/test_study_functions.py b/tests/test_study/test_study_functions.py index 7dc6b6d2a..e1decc443 100644 --- a/tests/test_study/test_study_functions.py +++ b/tests/test_study/test_study_functions.py @@ -262,3 +262,129 @@ def test_study_list(self): study_list = openml.study.list_studies(status="in_preparation") # might fail if server is recently reset assert len(study_list) >= 2 + + @pytest.mark.test_server() + def test_study_attach_runs_object_method(self): + run_list = openml.runs.list_runs(size=5) + assert len(run_list) == 5 + run_ids = list(run_list["run_id"]) + + study = openml.study.create_study( + alias=None, + benchmark_suite=None, + name="unit tested study attach runs", + description="test attach_runs", + run_ids=run_ids, + ) + study.publish() + TestBase._mark_entity_for_removal("study", study.id) + TestBase.logger.info(f"collected from {__file__.split('/')[-1]}: {study.id}") + + study_downloaded = openml.study.get_study(study.id) + self.assertSetEqual(set(study_downloaded.runs), set(run_ids)) + + # attach more runs using the object method + run_list_additional = openml.runs.list_runs(size=3, offset=5) + run_list_additional_ids = list(run_list_additional["run_id"]) + attached_count = study.attach_runs(run_list_additional_ids) + assert attached_count == len(run_ids) + len(run_list_additional_ids) + + # verify local state updated + self.assertSetEqual(set(study.runs), set(run_ids) | set(run_list_additional_ids)) + + study_downloaded = openml.study.get_study(study.id) + self.assertSetEqual(set(study_downloaded.runs), set(run_ids) | set(run_list_additional_ids)) + + # detach runs using the object method + detached_count = study.detach_runs(run_ids) + assert detached_count == len(run_list_additional_ids) + + # verify local state updated + self.assertSetEqual(set(study.runs), set(run_list_additional_ids)) + + study_downloaded = openml.study.get_study(study.id) + self.assertSetEqual(set(study_downloaded.runs), set(run_list_additional_ids)) + + @pytest.mark.test_server() + def test_study_attach_runs_unpublished_raises(self): + study = openml.study.create_study( + alias=None, + benchmark_suite=None, + name="unpublished study", + description="none", + run_ids=None, + ) + with pytest.raises(ValueError, match="Cannot attach runs to an unpublished study"): + study.attach_runs([1]) + + @pytest.mark.test_server() + def test_study_detach_runs_unpublished_raises(self): + study = openml.study.create_study( + alias=None, + benchmark_suite=None, + name="unpublished study", + description="none", + run_ids=None, + ) + with pytest.raises(ValueError, match="Cannot detach runs from an unpublished study"): + study.detach_runs([1]) + + @pytest.mark.test_server() + def test_suite_attach_tasks_object_method(self): + fixture_task_ids = [1, 2, 3] + + suite = openml.study.create_benchmark_suite( + alias=None, + name="unit tested suite attach tasks", + description="test attach_tasks", + task_ids=fixture_task_ids, + ) + suite.publish() + TestBase._mark_entity_for_removal("study", suite.id) + TestBase.logger.info(f"collected from {__file__.split('/')[-1]}: {suite.id}") + + suite_downloaded = openml.study.get_suite(suite.id) + self.assertSetEqual(set(suite_downloaded.tasks), set(fixture_task_ids)) + + # attach more tasks using the object method + tasks_additional = [4, 5, 6] + attached_count = suite.attach_tasks(tasks_additional) + assert attached_count == len(fixture_task_ids) + len(tasks_additional) + + # verify local state updated + self.assertSetEqual(set(suite.tasks), set(fixture_task_ids + tasks_additional)) + + suite_downloaded = openml.study.get_suite(suite.id) + self.assertSetEqual(set(suite_downloaded.tasks), set(fixture_task_ids + tasks_additional)) + + # detach tasks using the object method + detached_count = suite.detach_tasks(fixture_task_ids) + assert detached_count == len(tasks_additional) + + # verify local state updated + self.assertSetEqual(set(suite.tasks), set(tasks_additional)) + + suite_downloaded = openml.study.get_suite(suite.id) + self.assertSetEqual(set(suite_downloaded.tasks), set(tasks_additional)) + + @pytest.mark.test_server() + def test_suite_attach_tasks_unpublished_raises(self): + suite = openml.study.create_benchmark_suite( + alias=None, + name="unpublished suite", + description="none", + task_ids=[1], + ) + with pytest.raises(ValueError, match="Cannot attach tasks to an unpublished suite"): + suite.attach_tasks([2]) + + @pytest.mark.test_server() + def test_suite_detach_tasks_unpublished_raises(self): + suite = openml.study.create_benchmark_suite( + alias=None, + name="unpublished suite", + description="none", + task_ids=[1], + ) + with pytest.raises(ValueError, match="Cannot detach tasks from an unpublished suite"): + suite.detach_tasks([1])