diff --git a/openml/setups/setup.py b/openml/setups/setup.py index 0960ad4c1..170838138 100644 --- a/openml/setups/setup.py +++ b/openml/setups/setup.py @@ -1,12 +1,14 @@ # License: BSD 3-Clause from __future__ import annotations +from dataclasses import asdict, dataclass from typing import Any import openml.config import openml.flows +@dataclass class OpenMLSetup: """Setup object (a.k.a. Configuration). @@ -20,20 +22,20 @@ class OpenMLSetup: The setting of the parameters """ - def __init__(self, setup_id: int, flow_id: int, parameters: dict[int, Any] | None): - if not isinstance(setup_id, int): + setup_id: int + flow_id: int + parameters: dict[int, Any] | None + + def __post_init__(self) -> None: + if not isinstance(self.setup_id, int): raise ValueError("setup id should be int") - if not isinstance(flow_id, int): + if not isinstance(self.flow_id, int): raise ValueError("flow id should be int") - if parameters is not None and not isinstance(parameters, dict): + if self.parameters is not None and not isinstance(self.parameters, dict): raise ValueError("parameters should be dict") - self.setup_id = setup_id - self.flow_id = flow_id - self.parameters = parameters - def _to_dict(self) -> dict[str, Any]: return { "setup_id": self.setup_id, @@ -66,6 +68,7 @@ def __repr__(self) -> str: return header + body +@dataclass class OpenMLParameter: """Parameter object (used in setup). @@ -91,37 +94,24 @@ class OpenMLParameter: If the parameter was set, the value that it was set to. """ - def __init__( # noqa: PLR0913 - self, - input_id: int, - flow_id: int, - flow_name: str, - full_name: str, - parameter_name: str, - data_type: str, - default_value: str, - value: str, - ): - self.id = input_id - self.flow_id = flow_id - self.flow_name = flow_name - self.full_name = full_name - self.parameter_name = parameter_name - self.data_type = data_type - self.default_value = default_value - self.value = value + input_id: int + flow_id: int + flow_name: str + full_name: str + parameter_name: str + data_type: str + default_value: str + value: str + + def __post_init__(self) -> None: + # Map input_id to id for backward compatibility + self.id = self.input_id def _to_dict(self) -> dict[str, Any]: - return { - "id": self.id, - "flow_id": self.flow_id, - "flow_name": self.flow_name, - "full_name": self.full_name, - "parameter_name": self.parameter_name, - "data_type": self.data_type, - "default_value": self.default_value, - "value": self.value, - } + result = asdict(self) + # Replaces input_id with id for backward compatibility + result["id"] = result.pop("input_id") + return result def __repr__(self) -> str: header = "OpenML Parameter"