Merge pull request #179 from AAFC-MBB/master

Add Dataset Collections to Object Oriented API

bioblend/galaxy/objects/wrappers.py

@@ -26,6 +26,8 @@
     'Folder',
     'Dataset',
     'HistoryDatasetAssociation',
+    'DatasetCollection',
+    'HistoryDatasetCollectionAssociation',
     'LibraryDatasetDatasetAssociation',
     'LibraryDataset',
     'Tool',
@@ -445,10 +447,17 @@ def run(self, input_map=None, history='', params=None, import_inputs=False,
             raise TypeError(
                 'history must be either a history wrapper or a string')
         res = self.gi.gi.workflows.run_workflow(self.id, **kwargs)
-        # res structure: {'history': HIST_ID, 'outputs': [DS_ID, DS_ID, ...]}
+        # res structure: {'history': HIST_ID, 'outputs': [CI_ID, CI_ID, ...]}
         out_hist = self.gi.histories.get(res['history'])
-        assert set(res['outputs']).issubset(out_hist.dataset_ids)
-        outputs = [out_hist.get_dataset(_) for _ in res['outputs']]
+        content_infos_dict = dict()
+        for ci in out_hist.content_infos:
+            content_infos_dict[ci.id] = ci
+        outputs = []
+        for output_id in res['outputs']:
+            if content_infos_dict[output_id].type == 'file':
+                outputs.append(out_hist.get_dataset(output_id))
+            elif content_infos_dict[output_id].type == 'collection':
+                outputs.append(out_hist.get_dataset_collection(output_id))
 
         if wait:
             self.gi._wait_datasets(outputs, polling_interval=polling_interval,
@@ -626,6 +635,56 @@ def delete(self):
         self.refresh()
 
 
+@six.add_metaclass(abc.ABCMeta)
+class DatasetCollection(Wrapper):
+    """
+    Abstract base class for Galaxy dataset collections.
+    """
+    BASE_ATTRS = Wrapper.BASE_ATTRS + (
+        'state', 'deleted', 'collection_type'
+    )
+
+    @abc.abstractmethod
+    def __init__(self, dsc_dict, container, gi=None):
+        super(DatasetCollection, self).__init__(dsc_dict, gi=gi)
+        object.__setattr__(self, 'container', container)
+
+    def refresh(self):
+        """
+        Re-fetch the attributes pertaining to this object.
+
+        Returns: self
+        """
+        gi_client = getattr(self.gi.gi, self.container.API_MODULE)
+        dsc_dict = gi_client.show_dataset_collection(self.container.id, self.id)
+        self.__init__(dsc_dict, self.container, self.gi)
+        return self
+
+
+class HistoryDatasetCollectionAssociation(DatasetCollection):
+    """
+    Maps to a Galaxy ``HistoryDatasetCollectionAssociation``.
+    """
+    BASE_ATTRS = DatasetCollection.BASE_ATTRS + ('tags', 'visible', 'elements')
+    SRC = 'hdca'
+
+    def __init__(self, dsc_dict, container, gi=None):
+        super(HistoryDatasetCollectionAssociation, self).__init__(
+            dsc_dict, container, gi=gi)
+
+    @property
+    def gi_module(self):
+        return self.gi.histories
+
+    def delete(self):
+        """
+        Delete this dataset collection.
+        """
+        self.gi.gi.histories.delete_dataset_collection(self.container.id, self.id)
+        self.container.refresh()
+        self.refresh()
+
+
 class LibRelatedDataset(Dataset):
     """
     Base class for LibraryDatasetDatasetAssociation and LibraryDataset classes.
@@ -807,6 +866,7 @@ class History(DatasetContainer):
     """
     BASE_ATTRS = DatasetContainer.BASE_ATTRS + ('annotation', 'state', 'state_ids', 'state_details', 'tags')
     DS_TYPE = HistoryDatasetAssociation
+    DSC_TYPE = HistoryDatasetCollectionAssociation
     CONTENT_INFO_TYPE = HistoryContentInfo
     API_MODULE = 'histories'
 
@@ -965,6 +1025,33 @@ def download(self, jeha_id, outf, chunk_size=bioblend.CHUNK_SIZE):
         return self.gi.gi.histories.download_history(
             self.id, jeha_id, outf, chunk_size=chunk_size)
 
+    def create_dataset_collection(self, collection_description):
+        """
+        Create a new dataset collection in the history by providing the dataset ids.
+
+        :type collection_description: bioblend.galaxy.dataset_collections.CollectionDescription
+        :param collection_description: a description of the dataset collection
+
+        :rtype: class:`~.HistoryDatasetCollectionAssociation`
+        :return: the new dataset collection
+        """
+        dataset_collection = self.gi.gi.histories.create_dataset_collection(self.id, collection_description)
+        self.refresh()
+        return self.get_dataset_collection(dataset_collection['id'])
+
+    def get_dataset_collection(self, dsc_id):
+        """
+        Retrieve the dataset collection corresponding to the given id.
+
+        :type dsc_id: str
+        :param dsc_id: dataset collection id
+
+        :rtype: :class:`~.HistoryDatasetCollectionAssociation`
+        :return: the dataset collection corresponding to ``dsc_id``
+        """
+        dsc_dict = self.gi.gi.histories.show_dataset_collection(self.id, dsc_id)
+        return self.DSC_TYPE(dsc_dict, self, gi=self.gi)
+
 
 class Library(DatasetContainer):
     """

tests/TestGalaxyObjects.py

@@ -16,6 +16,7 @@
 import bioblend.galaxy.objects.wrappers as wrappers
 import bioblend.galaxy.objects.galaxy_instance as galaxy_instance
 from bioblend import ConnectionError
+from bioblend.galaxy import dataset_collections
 
 import test_util
 from test_util import unittest
@@ -619,6 +620,35 @@ def test_update(self):
         self.assertEqual(self.hist.annotation, new_annotation)
         self.assertEqual(self.hist.tags, new_tags)
 
+    @test_util.skip_unless_galaxy('release_14.06')
+    def test_create_dataset_collection(self):
+        self._create_collection_description()
+        hdca = self.hist.create_dataset_collection(self.collection_description)
+        self.assertIsInstance(hdca, wrappers.HistoryDatasetCollectionAssociation)
+        self.assertEqual(hdca.collection_type, 'list')
+        self.assertIs(hdca.container, self.hist)
+        self.assertEqual(len(hdca.elements), 2)
+        self.assertEqual(self.dataset1.id, hdca.elements[0]['object']['id'])
+        self.assertEqual(self.dataset2.id, hdca.elements[1]['object']['id'])
+
+    @test_util.skip_unless_galaxy('release_14.06')
+    def test_delete_dataset_collection(self):
+        self._create_collection_description()
+        hdca = self.hist.create_dataset_collection(self.collection_description)
+        hdca.delete()
+        self.assertTrue(hdca.deleted)
+
+    def _create_collection_description(self):
+        self.dataset1 = self.hist.paste_content(FOO_DATA)
+        self.dataset2 = self.hist.paste_content(FOO_DATA_2)
+        self.collection_description = dataset_collections.CollectionDescription(
+            name="MyDatasetList",
+            elements=[
+                dataset_collections.HistoryDatasetElement(name="sample1", id=self.dataset1.id),
+                dataset_collections.HistoryDatasetElement(name="sample2", id=self.dataset2.id),
+            ]
+        )
+
 
 @test_util.skip_unless_galaxy()
 class TestHDAContents(GalaxyObjectsTestBase):
@@ -706,6 +736,41 @@ def test_params(self):
         self.__test(params=True)
 
 
+@test_util.skip_unless_galaxy('release_14.08')
+class TestRunDatasetCollectionWorkflow(GalaxyObjectsTestBase):
+
+    def setUp(self):
+        super(TestRunDatasetCollectionWorkflow, self).setUp()
+        with open(SAMPLE_WF_COLL_FN) as f:
+            self.wf = self.gi.workflows.import_new(f.read())
+        self.hist = self.gi.histories.create('test_%s' % uuid.uuid4().hex)
+
+    def tearDown(self):
+        self.wf.delete()
+        self.hist.delete(purge=True)
+
+    def test_run_workflow_with_dataset_collection(self):
+        dataset1 = self.hist.paste_content(FOO_DATA)
+        dataset2 = self.hist.paste_content(FOO_DATA_2)
+        collection_description = dataset_collections.CollectionDescription(
+            name="MyDatasetList",
+            elements=[
+                dataset_collections.HistoryDatasetElement(name="sample1", id=dataset1.id),
+                dataset_collections.HistoryDatasetElement(name="sample2", id=dataset2.id),
+            ]
+        )
+        dataset_collection = self.hist.create_dataset_collection(collection_description)
+        input_map = {"Input Dataset Collection": dataset_collection,
+                     "Input 2": dataset1}
+        outputs, out_hist = self.wf.run(input_map, self.hist, wait=True)
+        self.assertEqual(len(outputs), 1)
+        out_hdca = outputs[0]
+        self.assertIsInstance(out_hdca, wrappers.HistoryDatasetCollectionAssociation)
+        self.assertEqual(out_hdca.collection_type, 'list')
+        self.assertEqual(len(out_hdca.elements), 2)
+        self.assertEqual(out_hist.id, self.hist.id)
+
+
 @test_util.skip_unless_galaxy()
 class TestJob(GalaxyObjectsTestBase):