FIX allow to retrieve class labels by the tasks class label flag

Author: mfeurer

openml/datasets/dataset.py

@@ -239,7 +239,7 @@ def get_data(self, target=None, target_dtype=int, include_row_id=False,
         else:
             return rval
 
-    def _retrieve_class_labels(self):
+    def _retrieve_class_labels(self, target_attribute='class'):
         """Reads the datasets arff to determine the class-labels, and returns those.
         If the task has no class labels (for example a regression problem) it returns None."""
         # TODO improve performance, currently reads the whole file
@@ -250,10 +250,8 @@ def _retrieve_class_labels(self):
             arffData = arff.ArffDecoder().decode(fh)
 
         dataAttributes = dict(arffData['attributes'])
-        if('class' in dataAttributes):
-            return dataAttributes['class']
-        elif('Class' in dataAttributes):
-            return dataAttributes['Class']
+        if target_attribute in dataAttributes:
+            return dataAttributes[target_attribute]
         else:
             return None
 

openml/tasks/functions.py

@@ -236,7 +236,7 @@ def get_task(task_id):
 
     # TODO look into either adding the class labels to task xml, or other
     # way of reading it.
-    class_labels = dataset._retrieve_class_labels()
+    class_labels = dataset._retrieve_class_labels(task.target_feature)
     task.class_labels = class_labels
     return task
 

tests/datasets/test_datasets.py

@@ -200,6 +200,15 @@ def test_publish_dataset(self):
         return_code, return_value = dataset.publish()
         self.assertEqual(return_code, 200)
 
+    def test__retrieve_class_labels(self):
+        openml.config.set_cache_directory(self.static_cache_dir)
+        labels = openml.datasets.get_dataset(2)._retrieve_class_labels()
+        self.assertEqual(labels, ['1', '2', '3', '4', '5', 'U'])
+        labels = openml.datasets.get_dataset(2)._retrieve_class_labels(
+            target_attribute='product-type')
+        self.assertEqual(labels, ['C', 'H', 'G'])
+        print(labels)
+
     def test_upload_dataset_with_url(self):
         dataset = OpenMLDataset(
             name="UploadTestWithURL", version=1, description="test",