fix test & pep8 & mypy

Author: mfeurer

openml/extensions/extension_interface.py

@@ -10,7 +10,7 @@
 if TYPE_CHECKING:
     from openml.flows import OpenMLFlow
     from openml.tasks.task import OpenMLTask
-    from openml.runs.trace import OpenMLRunTrace, OpenMLTraceIteration
+    from openml.runs.trace import OpenMLRunTrace, OpenMLTraceIteration  # noqa F401
 
 
 class Extension(ABC):

openml/extensions/sklearn/extension.py

@@ -1268,8 +1268,8 @@ def _prediction_to_probabilities(
 
             if proba_y.shape[1] != len(task.class_labels):
                 message = "Estimator only predicted for {}/{} classes!".format(
-                        proba_y.shape[1], len(task.class_labels),
-                    )
+                    proba_y.shape[1], len(task.class_labels),
+                )
                 warnings.warn(message)
                 openml.config.logger.warn(message)
 
@@ -1284,7 +1284,7 @@ def _prediction_to_probabilities(
 
         if self._is_hpo_class(model_copy):
             trace_data = self._extract_trace_data(model_copy, rep_no, fold_no)
-            trace = self._obtain_arff_trace(model_copy, trace_data)
+            trace = self._obtain_arff_trace(model_copy, trace_data)  # type: Optional[OpenMLRunTrace]  # noqa E501
         else:
             trace = None
 

openml/runs/functions.py

@@ -4,7 +4,6 @@
 from typing import Any, List, Optional, Set, Tuple, Union, TYPE_CHECKING  # noqa F401
 import warnings
 
-import numpy as np
 import sklearn.metrics
 import xmltodict
 
@@ -382,7 +381,6 @@ def _run_task_get_arffcontent(
     'OrderedDict[str, OrderedDict]',
 ]:
     arff_datacontent = []  # type: List[List]
-    arff_tracecontent = []  # type: List[List]
     traces = []  # type: List[OpenMLRunTrace]
     # stores fold-based evaluation measures. In case of a sample based task,
     # this information is multiple times overwritten, but due to the ordering

openml/runs/trace.py

@@ -1,7 +1,7 @@
 from collections import OrderedDict
 import json
 import os
-from typing import List
+from typing import List, Tuple  # noqa F401
 
 import arff
 import xmltodict
@@ -346,23 +346,22 @@ def trace_from_xml(cls, xml):
             )
             trace[(repeat, fold, iteration)] = current
 
-        return cls(None, trace)
+        return cls(run_id, trace)
 
     @classmethod
     def merge_traces(cls, traces: List['OpenMLRunTrace']):
         for i in range(1, len(traces)):
             if traces[i] != traces[i - 1]:
                 raise ValueError('Cannot merge traces!')
 
-        merged_trace = OrderedDict()
+        merged_trace = OrderedDict()  # type: OrderedDict[Tuple[int, int, int], OpenMLTraceIteration]  # noqa E501
 
         for trace in traces:
             for iteration in trace:
                 merged_trace[(iteration.repeat, iteration.fold, iteration.iteration)] = iteration
 
         return cls(None, merged_trace)
 
-
     def __str__(self):
         return '[Run id: %d, %d trace iterations]' % (
             -1 if self.run_id is None else self.run_id,

tests/test_extensions/test_sklearn_extension/test_sklearn_extension.py

@@ -1264,7 +1264,7 @@ def test_run_model_on_fold(self):
         # TODO add some mocking here to actually test the innards of this function, too!
         res = self.extension._run_model_on_fold(
             clf, task, 0, 0, 0,
-            add_local_measures=True)
+        )
 
         arff_datacontent, arff_tracecontent, user_defined_measures, model = res
         # predictions

tests/test_runs/test_run_functions.py

@@ -4,6 +4,7 @@
 import random
 import time
 import sys
+import unittest.mock
 
 import numpy as np
 
@@ -1052,8 +1053,11 @@ def test__run_task_get_arffcontent(self):
         num_folds = 10
         num_repeats = 1
 
+        flow = unittest.mock.Mock()
+        flow.name = 'dummy'
         clf = SGDClassifier(loss='log', random_state=1)
         res = openml.runs.functions._run_task_get_arffcontent(
+            flow=flow,
             extension=self.extension,
             model=clf,
             task=task,
@@ -1246,12 +1250,15 @@ def test_run_on_dataset_with_missing_labels(self):
         # labels only declared in the arff file, but is not present in the
         # actual data
 
+        flow = unittest.mock.Mock()
+        flow.name = 'dummy'
         task = openml.tasks.get_task(2)
 
         model = Pipeline(steps=[('Imputer', Imputer(strategy='median')),
                                 ('Estimator', DecisionTreeClassifier())])
 
         data_content, _, _, _ = _run_task_get_arffcontent(
+            flow=flow,
             model=model,
             task=task,
             extension=self.extension,
@@ -1267,6 +1274,8 @@ def test_run_on_dataset_with_missing_labels(self):
     def test_predict_proba_hardclassifier(self):
         # task 1 (test server) is important: it is a task with an unused class
         tasks = [1, 3, 115]
+        flow = unittest.mock.Mock()
+        flow.name = 'dummy'
 
         for task_id in tasks:
             task = openml.tasks.get_task(task_id)
@@ -1280,12 +1289,14 @@ def test_predict_proba_hardclassifier(self):
             ])
 
             arff_content1, _, _, _ = _run_task_get_arffcontent(
+                flow=flow,
                 model=clf1,
                 task=task,
                 extension=self.extension,
                 add_local_measures=True,
             )
             arff_content2, _, _, _ = _run_task_get_arffcontent(
+                flow=flow,
                 model=clf2,
                 task=task,
                 extension=self.extension,