Update trace.py to fix mypy error

openml · LennartPurucker · Oct 31, 2023 · Aug 16, 2023 · Aug 16, 2023 · Aug 16, 2023
commit e3c475a1c10ce68fdb5ebe86de539ed3ec74a7f2
diff --git a/openml/runs/trace.py b/openml/runs/trace.py
@@ -4,7 +4,7 @@
 from dataclasses import dataclass
 import json
 import os
-from typing import List, Tuple, Optional  # noqa F401
+from typing import List, Tuple, Optional, Dict, Union  # noqa F401
 
 import arff
 import xmltodict
@@ -19,6 +19,82 @@
 ]
 
 
+@dataclass
+class OpenMLTraceIteration:
+    """
+    OpenML Trace Iteration: parsed output from Run Trace call
+    Exactly one of `setup_string` or `parameters` must be provided.
+
+    Parameters
+    ----------
+    repeat : int
+        repeat number (in case of no repeats: 0)
+
+    fold : int
+        fold number (in case of no folds: 0)
+
+    iteration : int
+        iteration number of optimization procedure
+
+    setup_string : str, optional
+        json string representing the parameters
+        If not provided, ``parameters`` should be set.
+
+    evaluation : double
+        The evaluation that was awarded to this trace iteration.
+        Measure is defined by the task
+
+    selected : bool
+        Whether this was the best of all iterations, and hence
+        selected for making predictions. Per fold/repeat there
+        should be only one iteration selected
+
+    parameters : OrderedDict, optional
+        Dictionary specifying parameter names and their values.
+        If not provided, ``setup_string`` should be set.
+    """
+
+    repeat: int
+    fold: int
+    iteration: int
+
+    evaluation: float
+    selected: bool
+
+    setup_string: Optional[str] = None
+    parameters: Optional[OrderedDict] = None
+
+    def __post_init__(self):
+        # TODO: refactor into one argument of type <str | OrderedDict>
+        if self.setup_string and self.parameters:
+            raise ValueError(
+                "Can only be instantiated with either `setup_string` or `parameters` argument."
+            )
+        elif not (self.setup_string or self.parameters):
+            raise ValueError(
+                "Either `setup_string` or `parameters` needs to be passed as argument."
+            )
+        if self.parameters is not None and not isinstance(self.parameters, OrderedDict):
+            raise TypeError(
+                "argument parameters is not an instance of OrderedDict, but %s"
+                % str(type(self.parameters))
+            )
+
+    def get_parameters(self):
+        result = {}
+        # parameters have prefix 'parameter_'
+
+        if self.setup_string:
+            for param in self.setup_string:
+                key = param[len(PREFIX) :]
+                value = self.setup_string[param]
+                result[key] = json.loads(value)
+        else:
+            for param, value in self.parameters.items():
+                result[param[len(PREFIX) :]] = value
+        return result
+
+
 class OpenMLRunTrace(object):
     """OpenML Run Trace: parsed output from Run Trace call
 
@@ -33,7 +109,11 @@ class OpenMLRunTrace(object):
 
     """
 
-    def __init__(self, run_id: int, trace_iterations: List[List]):
+    def __init__(
+        self,
+        run_id: Union[int, None],
+        trace_iterations: Dict[Tuple[int, int, int], OpenMLTraceIteration],
+    ):
         """Object to hold the trace content of a run.
 
         Parameters
@@ -431,79 +511,3 @@ def __repr__(self):
     def __iter__(self):
         for val in self.trace_iterations.values():
             yield val
-
-
-@dataclass
-class OpenMLTraceIteration:
-    """
-    OpenML Trace Iteration: parsed output from Run Trace call
-    Exactly one of `setup_string` or `parameters` must be provided.
-
-    Parameters
-    ----------
-    repeat : int
-        repeat number (in case of no repeats: 0)
-
-    fold : int
-        fold number (in case of no folds: 0)
-
-    iteration : int
-        iteration number of optimization procedure
-
-    setup_string : str, optional
-        json string representing the parameters
-        If not provided, ``parameters`` should be set.
-
-    evaluation : double
-        The evaluation that was awarded to this trace iteration.
-        Measure is defined by the task
-
-    selected : bool
-        Whether this was the best of all iterations, and hence
-        selected for making predictions. Per fold/repeat there
-        should be only one iteration selected
-
-    parameters : OrderedDict, optional
-        Dictionary specifying parameter names and their values.
-        If not provided, ``setup_string`` should be set.
-    """
-
-    repeat: int
-    fold: int
-    iteration: int
-
-    evaluation: float
-    selected: bool
-
-    setup_string: Optional[str] = None
-    parameters: Optional[OrderedDict] = None
-
-    def __post_init__(self):
-        # TODO: refactor into one argument of type <str | OrderedDict>
-        if self.setup_string and self.parameters:
-            raise ValueError(
-                "Can only be instantiated with either `setup_string` or `parameters` argument."
-            )
-        elif not (self.setup_string or self.parameters):
-            raise ValueError(
-                "Either `setup_string` or `parameters` needs to be passed as argument."
-            )
-        if self.parameters is not None and not isinstance(self.parameters, OrderedDict):
-            raise TypeError(
-                "argument parameters is not an instance of OrderedDict, but %s"
-                % str(type(self.parameters))
-            )
-
-    def get_parameters(self):
-        result = {}
-        # parameters have prefix 'parameter_'
-
-        if self.setup_string:
-            for param in self.setup_string:
-                key = param[len(PREFIX) :]
-                value = self.setup_string[param]
-                result[key] = json.loads(value)
-        else:
-            for param, value in self.parameters.items():
-                result[param[len(PREFIX) :]] = value
-        return result