From 33803db2cb9befbfd5fef11f81a8bf2ebd014185 Mon Sep 17 00:00:00 2001 From: arcwiser Date: Tue, 13 Oct 2026 01:07:15 +1100 Subject: [PATCH] Change list_setups default to dataframe Signed-off-by: arcwiser --- openml/setups/functions.py | 10 +++++-- tests/test_setups/test_setup_functions.py | 35 ++++++++++++++++------- 2 files changed, 32 insertions(+), 13 deletions(-) diff --git a/openml/setups/functions.py b/openml/setups/functions.py index 0b1a4f375..13bd2c6d9 100644 --- a/openml/setups/functions.py +++ b/openml/setups/functions.py @@ -77,7 +77,7 @@ def list_setups( # noqa: PLR0913 flow: int | None = None, tag: str | None = None, setup: Iterable[int] | None = None, - output_format: Literal["object", "dataframe"] = "object", + output_format: Literal["object", "dataframe"] = "dataframe", ) -> dict[int, OpenMLSetup] | pd.DataFrame: """ List all setups matching all of the given filters. @@ -89,7 +89,7 @@ def list_setups( # noqa: PLR0913 flow : int, optional tag : str, optional setup : Iterable[int], optional - output_format: str, optional (default='object') + output_format: str, optional (default='dataframe') The parameter decides the format of the output. - If 'dataframe' the output is a pandas DataFrame - If 'object' the output is a dictionary of OpenMLSetup objects @@ -115,7 +115,11 @@ def list_setups( # noqa: PLR0913 return {setup.setup_id: setup for setup in flattened} records = [setup._to_dict() for setup in flattened] - return pd.DataFrame.from_records(records, index="setup_id") + return pd.DataFrame.from_records( + records, + columns=["setup_id", "flow_id", "parameters"], + index="setup_id", + ) def initialize_model(setup_id: int, *, strict_version: bool = True) -> Any: diff --git a/tests/test_setups/test_setup_functions.py b/tests/test_setups/test_setup_functions.py index 0735925f2..b0f18c12a 100644 --- a/tests/test_setups/test_setup_functions.py +++ b/tests/test_setups/test_setup_functions.py @@ -17,6 +17,21 @@ from openml.testing import TestBase +@pytest.mark.parametrize( + ("batches", "expected_index"), + [([[openml.setups.setup.OpenMLSetup(7, 11, None)]], [7]), ([], [])], +) +def test_list_setups_defaults_to_dataframe(mocker, batches, expected_index): + mocker.patch("openml.utils._list_all", return_value=batches) + + setups = openml.setups.list_setups() + + assert isinstance(setups, pd.DataFrame) + assert setups.index.name == "setup_id" + assert list(setups.index) == expected_index + assert list(setups.columns) == ["flow_id", "parameters"] + + def get_sentinel(): # Create a unique prefix for the flow. Necessary because the flow is # identified by its name and external version online. Having a unique @@ -141,7 +156,7 @@ def test_setup_list_filter_flow(self): flow_id = 5873 - setups = openml.setups.list_setups(flow=flow_id) + setups = openml.setups.list_setups(flow=flow_id, output_format="object") assert len(setups) >= 2 for setup_id in setups: @@ -150,22 +165,22 @@ def test_setup_list_filter_flow(self): @pytest.mark.test_server() def test_list_setups_empty(self): setups = openml.setups.list_setups(setup=[0]) - if len(setups) > 0: - raise ValueError("UnitTest Outdated, got somehow results") - - assert isinstance(setups, dict) + assert isinstance(setups, pd.DataFrame) + assert setups.empty + assert setups.index.name == "setup_id" + assert list(setups.columns) == ["flow_id", "parameters"] @pytest.mark.production_server() def test_list_setups_output_format(self): self.use_production_server() flow_id = 6794 setups = openml.setups.list_setups(flow=flow_id, size=10) - assert isinstance(setups, dict) - assert isinstance(setups[next(iter(setups.keys()))], openml.setups.setup.OpenMLSetup) + assert isinstance(setups, pd.DataFrame) assert len(setups) == 10 - setups = openml.setups.list_setups(flow=flow_id, size=10, output_format="dataframe") - assert isinstance(setups, pd.DataFrame) + setups = openml.setups.list_setups(flow=flow_id, size=10, output_format="object") + assert isinstance(setups, dict) + assert isinstance(setups[next(iter(setups))], openml.setups.setup.OpenMLSetup) assert len(setups) == 10 @pytest.mark.test_server() @@ -176,7 +191,7 @@ def test_setuplist_offset(self): setups2 = openml.setups.list_setups(offset=size, size=size) assert len(setups2) == size - all = set(setups.keys()).union(setups2.keys()) + all = set(setups.index).union(setups2.index) assert len(all) == size * 2