Merge pull request planetlabs#71 from planetlabs/sync-more-products

sync tool handles all ortho products

Author: ischneider

planet/api/sync.py

@@ -24,11 +24,13 @@
 
 class _SyncTool(object):
 
-    def __init__(self, client, destination, aoi, scene_type, **filters):
+    def __init__(self, client, destination, aoi, scene_type, products,
+                 **filters):
         self.client = client
         self.destination = destination
         self.aoi = aoi
         self.scene_type = scene_type
+        self.products = products
         self.filters = filters
         self.workspace = filters.get('workspace', None)
         self._init()
@@ -75,13 +77,14 @@ def init(self, limit=-1):
         self._scenes = resp
         count = resp.get()['count']
         self._scene_count = count if limit < 0 else limit
-        return count
+        return count * len(self.products)
 
     def get_scenes_to_sync(self):
         return self._scenes.items_iter(limit=self._scene_count)
 
     def sync(self, callback):
-        summary = _SyncSummary(self._scene_count)
+        # init with number of scenes * products
+        summary = _SyncSummary(self._scene_count * len(self.products))
 
         all_scenes = self.get_scenes_to_sync()
         while True:
@@ -93,9 +96,12 @@ def sync(self, callback):
                 _SyncHandler(self.destination, summary, scene, callback) for
                 scene in scenes
             ]
-            futures = [h.run(self.client, self.scene_type) for h in handlers]
-            for f in futures:
-                f.await()
+            # start all downloads asynchronously
+            for h in handlers:
+                h.run(self.client, self.scene_type, self.products)
+            # synchronously await them and then write metadata
+            for h in handlers:
+                h.finish()
 
         if summary.latest:
             sync = self._read_sync_file()
@@ -132,23 +138,29 @@ def __init__(self, destination, summary, metadata, user_callback):
         self.metadata = metadata
         self.user_callback = user_callback or (lambda *args: None)
 
-    def run(self, client, scene_type):
-        return client.fetch_scene_geotiffs(
-            [self.metadata['id']],
-            scene_type,
-            callback=self)[0]
+    def run(self, client, scene_type, products):
+        self.futures = []
+        for product in products:
+            self.futures.extend(client.fetch_scene_geotiffs(
+                                [self.metadata['id']],
+                                scene_type, product,
+                                callback=self))
 
-    def __call__(self, body):
-        '''implement the callback that runs when the scene response is ready'''
-        # first write scene to disk
-        write_to_file(self.destination)(body)
+    def finish(self):
+        for f in self.futures:
+            f.await()
 
         # write out metadata
         metadata = os.path.join(self.destination,
                                 '%s_metadata.json' % self.metadata['id'])
         with atomic_open(metadata, 'wb') as fp:
             fp.write(json.dumps(self.metadata, indent=2).encode('utf-8'))
 
+    def __call__(self, body):
+        '''implement the callback that runs when the scene response is ready'''
+        # first write scene to disk
+        write_to_file(self.destination)(body)
+
         # summarize
         self.summary.transfer_complete(body, self.metadata)
 

planet/scripts/__init__.py

@@ -30,6 +30,8 @@
 
 client_params = {}
 
+ORTHO_PRODUCTS = ['visual', 'analytic', 'unrectified']
+
 
 def client():
     return api.Client(**client_params)
@@ -268,7 +270,7 @@ def get_scenes_list(scene_type, pretty, aoi, limit, where, workspace):
 
     echo_json_response(call_and_wrap(
         client().get_scenes_list,
-        scene_type=scene_type, intersects=aoi,
+        scene_type=scene_type, intersects=aoi, count=limit,
         **conditions), pretty, limit=limit)
 
 
@@ -289,7 +291,7 @@ def metadata(scene_id, scene_type, pretty):
 @click.option('--product',
               type=click.Choice(
                   ["band_%d" % i for i in range(1, 12)] +
-                  ['visual', 'analytic', 'unrectified', 'qa']
+                  ORTHO_PRODUCTS + ['qa']
               ), default='visual')
 @cli.command('download')
 def fetch_scene_geotiff(scene_ids, scene_type, product, dest):
@@ -333,20 +335,28 @@ def fetch_scene_thumbnails(scene_ids, scene_type, size, fmt, dest):
 @click.argument("destination")
 @limit_option(default=-1)
 @click.option("--dryrun", is_flag=True, help='Do not actually download')
+@click.option("--products", multiple=True,
+              type=click.Choice(ORTHO_PRODUCTS + ['all']),
+              help='Specifiy products to download, default is visual')
 @cli.command('sync')
-def sync(destination, workspace, scene_type, limit, dryrun):
+def sync(destination, workspace, scene_type, limit, dryrun, products):
     '''Synchronize a directory to a specified AOI or workspace'''
     aoi = None
     filters = {'workspace': workspace}
 
+    if 'all' in products:
+        products = ORTHO_PRODUCTS
+    else:
+        products = products or ('visual',)
+
     sync_tool = _SyncTool(client(), destination, aoi,
-                          scene_type, **filters)
+                          scene_type, products, **filters)
 
     try:
         to_fetch = sync_tool.init(limit)
     except ValueError as ve:
         raise click.ClickException(str(ve))
-    click.echo('total scenes to fetch: %s' % to_fetch)
+    click.echo('total scene products to fetch: %s' % to_fetch)
     if limit > -1:
         click.echo('limiting to %s' % limit)
 

tests/test_sync.py

@@ -26,22 +26,22 @@
 def test_sync_tool(tmpdir):
     # test non-existing destination
     try:
-        _SyncTool(client, 'should-not-exist', None, None)
+        _SyncTool(client, 'should-not-exist', None, None, None)
     except ValueError as ve:
         assert str(ve) == 'destination must exist and be a directory'
 
     # test existing destination, no aoi.geojson
     td = tmpdir.mkdir('sync-dest')
     try:
-        _SyncTool(client, td.strpath, None, None)
+        _SyncTool(client, td.strpath, None, None, None)
     except ValueError as ve:
         assert str(ve) == 'no aoi provided and no aoi.geojson file'
 
     # test existing destination, invalid aoi.geojson
     aoi_file = td.join('aoi.geojson')
     aoi_file.write('not geojson')
     try:
-        _SyncTool(client, td.strpath, None, None)
+        _SyncTool(client, td.strpath, None, None, None)
     except ValueError as ve:
         assert str(ve) == '%s does not contain valid JSON' % aoi_file
 
@@ -51,19 +51,22 @@ def test_sync_tool(tmpdir):
 def test_sync_tool_init(tmpdir):
     td = tmpdir.mkdir('sync-dest')
     aoi = read_fixture('aoi.geojson')
-    st = _SyncTool(client, td.strpath, aoi, 'ortho')
+    st = _SyncTool(client, td.strpath, aoi, 'ortho', ('visual', 'analytic'))
     search_json = json.loads(read_fixture('search.geojson'))
     response = MagicMock(spec=Scenes)
     response.get.return_value = search_json
     client.get_scenes_list.return_value = response
 
     # init w/ no limit should return count from response
     count = st.init()
-    assert search_json['count'] == count
+    # confusing but we'll download 2 products one for each scene
+    assert search_json['count'] * 2 == count
 
-    # expect limiting
+    # expect limiting to 10 ids
     count = st.init(limit=10)
-    assert count == search_json['count']
+    # still replies with total jobs despite the limit
+    assert search_json['count'] * 2 == count
+    # this tracks the internal number of ids, still 10
     assert 10 == st._scene_count
 
     # create a stored 'latest' date and ensure it's used
@@ -82,7 +85,7 @@ def test_sync_tool_init(tmpdir):
 def test_sync_tool_sync(tmpdir):
     td = tmpdir.mkdir('sync-dest')
     aoi = read_fixture('aoi.geojson')
-    st = _SyncTool(client, td.strpath, aoi, 'ortho')
+    st = _SyncTool(client, td.strpath, aoi, 'ortho', ('visual',))
     search_json = json.loads(read_fixture('search.geojson'))
 
     class Page:
@@ -146,7 +149,7 @@ def callback(name, remaining):
 
     # because process is normally async, as a sideeffect of this mock getting
     # called, we have to dispatch the callback
-    def run_callbacks(ids, scene_type, callback):
+    def run_callbacks(ids, scene_type, product, callback):
         resp = responses.pop(0)
         callback(resp)
         return DEFAULT
@@ -172,7 +175,8 @@ def run_callbacks(ids, scene_type, callback):
     # implementation detail - because we're making requests separately,
     # fetch_scene_geotiffs will be called once for each id (instead of in bulk)
     assert items == len(args)
-    assert [a[0] for a in args] == [([f['id']], 'ortho') for f in first_items]
+    assert [a[0] for a in args] == \
+        [([f['id']], 'ortho', 'visual') for f in first_items]
     # callbacks should be made - arguments are 'tiff name', remaining
     assert called_back == [(str(i), 4 - i) for i in range(5)]