Merge pull request #668 from openml/fixes

Fix existing mypy issues

Author: PGijsbers

openml/datasets/dataset.py

@@ -1,9 +1,10 @@
+from collections import OrderedDict
 import gzip
 import io
 import logging
 import os
 import pickle
-from collections import OrderedDict
+from typing import List, Optional, Union
 
 import arff
 import numpy as np
@@ -417,7 +418,7 @@ def _download_data(self) -> None:
         from .functions import _get_dataset_arff
         self.data_file = _get_dataset_arff(self)
 
-    def get_data(self, target: str = None,
+    def get_data(self, target: Optional[Union[List[str], str]] = None,
                  include_row_id: bool = False,
                  include_ignore_attributes: bool = False,
                  return_categorical_indicator: bool = False,

openml/datasets/functions.py

@@ -375,10 +375,7 @@ def get_dataset(dataset_id: Union[int, str], download_data: bool = True) -> Open
             features = _get_dataset_features(did_cache_dir, dataset_id)
             qualities = _get_dataset_qualities(did_cache_dir, dataset_id)
 
-            if download_data:
-                arff_file = _get_dataset_arff(description)
-            else:
-                arff_file = None
+            arff_file = _get_dataset_arff(description) if download_data else None
 
             remove_dataset_cache = False
         except OpenMLServerException as e:

openml/flows/functions.py

@@ -79,7 +79,7 @@ def get_flow(flow_id: int, reinstantiate: bool = False) -> OpenMLFlow:
         The OpenML flow id.
 
     reinstantiate: bool
-        Whether to reinstantiate the flow to a sklearn model.
+        Whether to reinstantiate the flow to a model instance.
 
     Returns
     -------

tests/test_study/test_study_examples.py

@@ -26,6 +26,7 @@ def test_Figure1a(self):
             print('URL for run: %s/run/%d' %(openml.config.server,run.run_id))
         """  # noqa: E501
         import openml
+        import sklearn.metrics
         import sklearn.pipeline
         import sklearn.preprocessing
         import sklearn.tree