[ENH] Implement tagging for OpenMLSetup (openml#1686)

- Make OpenMLSetup inherit from OpenMLBase to enable tagging
- Add OpenMLSetup to _get_rest_api_type_alias mapping
- Add test_tagging test for push_tag and remove_tag

Author: ravencore06

openml/setups/setup.py

@@ -1,15 +1,15 @@
 # License: BSD 3-Clause
 from __future__ import annotations
 
+from collections.abc import Sequence
 from dataclasses import asdict, dataclass
 from typing import Any
 
-import openml.config
 import openml.flows
+from openml.base import OpenMLBase
 
 
-@dataclass
-class OpenMLSetup:
+class OpenMLSetup(OpenMLBase):
     """Setup object (a.k.a. Configuration).
 
     Parameters
@@ -22,33 +22,32 @@ class OpenMLSetup:
         The setting of the parameters
     """
 
-    setup_id: int
-    flow_id: int
-    parameters: dict[int, Any] | None
-
-    def __post_init__(self) -> None:
-        if not isinstance(self.setup_id, int):
+    def __init__(
+        self,
+        setup_id: int,
+        flow_id: int,
+        parameters: dict[int, Any] | None,
+    ) -> None:
+        if not isinstance(setup_id, int):
             raise ValueError("setup id should be int")
 
-        if not isinstance(self.flow_id, int):
+        if not isinstance(flow_id, int):
             raise ValueError("flow id should be int")
 
-        if self.parameters is not None and not isinstance(self.parameters, dict):
+        if parameters is not None and not isinstance(parameters, dict):
             raise ValueError("parameters should be dict")
 
-    def _to_dict(self) -> dict[str, Any]:
-        return {
-            "setup_id": self.setup_id,
-            "flow_id": self.flow_id,
-            "parameters": {p.id: p._to_dict() for p in self.parameters.values()}
-            if self.parameters is not None
-            else None,
-        }
+        self.setup_id = setup_id
+        self.flow_id = flow_id
+        self.parameters = parameters
 
-    def __repr__(self) -> str:
-        header = "OpenML Setup"
-        header = f"{header}\n{'=' * len(header)}\n"
+    @property
+    def id(self) -> int | None:
+        """The ID of the setup."""
+        return self.setup_id
 
+    def _get_repr_body_fields(self) -> Sequence[tuple[str, str | int | list[str] | None]]:
+        """Collect all information to display in the __repr__ body."""
         fields = {
             "Setup ID": self.setup_id,
             "Flow ID": self.flow_id,
@@ -57,15 +56,21 @@ def __repr__(self) -> str:
                 len(self.parameters) if self.parameters is not None else float("nan")
             ),
         }
-
-        # determines the order in which the information will be printed
         order = ["Setup ID", "Flow ID", "Flow URL", "# of Parameters"]
-        _fields = [(key, fields[key]) for key in order if key in fields]
+        return [(key, fields[key]) for key in order if key in fields]
 
-        longest_field_name_length = max(len(name) for name, _ in _fields)
-        field_line_format = f"{{:.<{longest_field_name_length}}}: {{}}"
-        body = "\n".join(field_line_format.format(name, value) for name, value in _fields)
-        return header + body
+    def _to_dict(self) -> dict[str, Any]:
+        return {
+            "setup_id": self.setup_id,
+            "flow_id": self.flow_id,
+            "parameters": {p.id: p._to_dict() for p in self.parameters.values()}
+            if self.parameters is not None
+            else None,
+        }
+
+    def _parse_publish_response(self, xml_response: dict[str, str]) -> None:
+        msg = "Setups cannot be published directly."
+        raise NotImplementedError(msg)
 
 
 @dataclass

openml/utils/_openml.py

@@ -99,6 +99,7 @@ def _get_rest_api_type_alias(oml_object: OpenMLBase) -> str:
     rest_api_mapping: list[tuple[type | tuple, str]] = [
         (openml.datasets.OpenMLDataset, "data"),
         (openml.flows.OpenMLFlow, "flow"),
+        (openml.setups.OpenMLSetup, "setup"),
         (openml.tasks.OpenMLTask, "task"),
         (openml.runs.OpenMLRun, "run"),
         ((openml.study.OpenMLStudy, openml.study.OpenMLBenchmarkSuite), "study"),

tests/test_setups/test_setup_functions.py

@@ -189,3 +189,20 @@ def test_get_uncached_setup(self):
         openml.config.set_root_cache_directory(self.static_cache_dir)
         with pytest.raises(openml.exceptions.OpenMLCacheException):
             openml.setups.functions._get_cached_setup(10)
+
+    @pytest.mark.test_server()
+    def test_tagging(self):
+        setups = openml.setups.list_setups(size=1)
+        setup_id = next(iter(setups.keys()))
+        setup = openml.setups.get_setup(setup_id)
+        unique_indicator = str(time.time()).replace(".", "")
+        tag = f"test_tag_TestSetup_{unique_indicator}"
+        tagged_setups = openml.setups.list_setups(tag=tag)
+        assert len(tagged_setups) == 0
+        setup.push_tag(tag)
+        tagged_setups = openml.setups.list_setups(tag=tag)
+        assert len(tagged_setups) == 1
+        assert setup_id in tagged_setups
+        setup.remove_tag(tag)
+        tagged_setups = openml.setups.list_setups(tag=tag)
+        assert len(tagged_setups) == 0