adds tagging and untagging functions

Author: janvanrijn

openml/utils.py

@@ -1,4 +1,6 @@
+import xmltodict
 import six
+from ._api_calls import _perform_api_call
 
 
 def extract_xml_tags(xml_tag_name, node, allow_none=True):
@@ -37,4 +39,50 @@ def extract_xml_tags(xml_tag_name, node, allow_none=True):
             return None
         else:
             raise ValueError("Could not find tag '%s' in node '%s'" %
-                             (xml_tag_name, str(node)))
+                             (xml_tag_name, str(node)))
+
+
+def _tag_entity(entity_type, entity_id, tag, untag=False):
+    """Abstract function that can be used as a partial for tagging entities
+       on OpenML
+
+        Parameters
+        ----------
+        entity_type : str
+            Name of the entity to tag (e.g., run, flow, data)
+
+        entity_id : int
+            OpenML id of the entity
+
+        tag : str
+            The tag
+
+        untag : bool
+            Set to true if needed to untag, rather than tag
+
+        Returns
+        -------
+        tags : list
+            List of tags that the entity is still tagged with
+        """
+    legal_entities = {'data', 'task', 'flow', 'setup', 'run'}
+    if entity_type not in legal_entities:
+        raise ValueError('Can\'t tag a %s' %entity_type)
+
+    uri = '%s/tag' %entity_type
+    main_tag = 'oml:%s_tag' %entity_type
+    if untag:
+        uri = '%s/untag' %entity_type
+        main_tag = 'oml:%s_untag' %entity_type
+
+
+    post_variables = {'%s_id'%entity_type: entity_id, 'tag': tag}
+    result_xml = _perform_api_call(uri, post_variables)
+
+    result = xmltodict.parse(result_xml, force_list={'oml:tag'})[main_tag]
+
+    if 'oml:tag' in result:
+        return result['oml:tag']
+    else:
+        # no tags, return empty list
+        return []

tests/test_datasets/test_dataset_functions.py

@@ -8,13 +8,15 @@
 else:
     import mock
 
+import random
 import six
 import scipy.sparse
 
 import openml
 from openml import OpenMLDataset
 from openml.exceptions import OpenMLCacheException, PyOpenMLError
 from openml.testing import TestBase
+from openml.utils import _tag_entity
 
 from openml.datasets.functions import (_get_cached_dataset,
                                        _get_cached_dataset_features,
@@ -105,6 +107,15 @@ def _check_dataset(self, dataset):
             self.assertIn(dataset['status'], ['in_preparation', 'active',
                                               'deactivated'])
 
+    def test_tag_untag_dataset(self):
+        tag = 'test_tag_%d' %random.randint(1, 1000000)
+        all_tags = _tag_entity('data', 1, tag)
+        self.assertTrue(tag in all_tags)
+        all_tags = _tag_entity('data', 1, tag, untag=True)
+        self.assertTrue(tag not in all_tags)
+
+
+
     def test_list_datasets(self):
         # We can only perform a smoke test here because we test on dynamic
         # data from the internet...