add download_assets functions + tests

Author: matthewhanson

satstac/item.py

@@ -123,6 +123,15 @@ def substitute(self, string):
                 subs[key] = self[key.replace('_colon_', ':')]
         return Template(string).substitute(**subs)   
 
+    def download_assets(self, keys=None, **kwargs):
+        """ Download multiple assets """
+        if keys is None:
+            keys = self.data['assets'].keys()
+        filenames = []
+        for key in keys:
+            filenames.append(self.download(key, **kwargs))
+        return filenames
+
     def download(self, key, overwrite=False, path='', filename='${id}', requestor_pays=False):
         """ Download this key (e.g., a band, or metadata file) from the scene """
         asset = self.asset(key)

satstac/items.py

@@ -118,6 +118,14 @@ def filter(self, key, values):
             items += list(filter(lambda x: x[key] == val, self._items))
         self._items = items
 
+    def download_assets(self, *args, **kwargs):
+        filenames = []
+        for i in self._items:
+            fnames = i.download_assets(*args, **kwargs)
+            if len(fnames) > 0:
+                filenames.append(fnames)
+        return filenames
+
     def download(self, *args, **kwargs):
         """ Download all Items """
         dls = []

test/test_item.py

@@ -97,6 +97,13 @@ def test_download(self):
         fname = item.download(key='MTL', path=self.path)
         assert(os.path.exists(fname))
 
+    def test_download_assets(self):
+        """ Retrieve multiple data files """
+        item = Item.open(self.filename)
+        fnames = item.download_assets(keys=['MTL', 'ANG'], path=self.path)
+        for f in fnames:
+            assert(os.path.exists(f))
+
     def test_download_nonexist(self):
         """ Test downloading of non-existent file """
         item = Item.open(self.filename)

test/test_items.py

@@ -2,12 +2,21 @@
 import unittest
 
 from satstac import Items, Item
+from shutil import rmtree
 
 testpath = os.path.dirname(__file__)
 
 
 class Test(unittest.TestCase):
 
+    path = os.path.join(testpath, 'test-item')
+
+    @classmethod
+    def tearDownClass(cls):
+        """ Remove test files """
+        if os.path.exists(cls.path):
+            rmtree(cls.path)
+
     def load_items(self):
         return Items.load(os.path.join(testpath, 'items.json'))
 
@@ -89,14 +98,21 @@ def test_filter(self):
         items.filter('eo:cloud_cover', [100])
         assert(len(items) == 1)
 
+    def test_download_assets(self):
+        """ Download multiple assets from all items """
+        items = self.load_items()
+        filenames = items.download_assets(keys=['MTL', 'ANG'], path=self.path)
+        assert(len(filenames) == 2)
+        for fnames in filenames:
+            assert(len(fnames) == 2)
+            for f in fnames:
+                assert(os.path.exists(f))
+
     def test_download(self):
         """ Download a data file from all items """
         items = self.load_items()
 
-        fnames = items.download(key='MTL')
+        fnames = items.download(key='MTL', path=self.path)
         assert(len(fnames) == 2)
         for f in fnames:
             assert(os.path.exists(f))
-            os.remove(f)
-            assert(not os.path.exists(f))
-        #shutil.rmtree(os.path.join(testpath, 'landsat-8-l1'))