Add attach_runs/detach_runs to OpenMLStudy and attach_tasks/detach_tasks to OpenMLBenchmarkSuite

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

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