Added method Database.download_from_vso_query_result

Makes it possible to not only add entries from a VSO query result but
also download the corresponding data files.

CHANGELOG.md

@@ -6,6 +6,7 @@ Latest
   - Added a window/split method to time range
   - Updates to spectrogram documentation
   - Added method Database.add_from_hek_query_result to HEK database
+  - Added method Database.download_from_vso_query_result
 
 0.4.0
 -----

sunpy/database/database.py

@@ -296,6 +296,32 @@ def commit(self):
         """
         self.session.commit()
 
+    def _download_and_collect_entries(self, query_result, client=None,
+            path=None, progress=False):
+        if client is None:
+            client = VSOClient()
+        for block in query_result:
+            paths = client.get([block], path).wait(progress=progress)
+            for path in paths:
+                qr_entry = tables.DatabaseEntry._from_query_result_block(block)
+                file_entries = list(
+                    tables.entries_from_file(path, self.default_waveunit))
+                for entry in file_entries:
+                    entry.source = qr_entry.source
+                    entry.provider = qr_entry.provider
+                    entry.physobs = qr_entry.physobs
+                    entry.fileid = qr_entry.fileid
+                    entry.observation_time_start =\
+                        qr_entry.observation_time_start
+                    entry.observation_time_end = qr_entry.observation_time_end
+                    entry.instrument = qr_entry.instrument
+                    entry.size = qr_entry.size
+                    entry.wavemin = qr_entry.wavemin
+                    entry.wavemax = qr_entry.wavemax
+                    entry.path = path
+                    entry.download_time = datetime.utcnow()
+                    yield entry
+
     def download(self, *query, **kwargs):
         """download(*query, client=sunpy.net.vso.VSOClient(), path=None, progress=False)
         Search for data using the VSO interface (see
@@ -323,28 +349,8 @@ def download(self, *query, **kwargs):
         # don't do anything if querying the VSO results in no data
         if not qr:
             return
-        entries = []
-        for block in qr:
-            paths = client.get([block], path).wait(progress=progress)
-            for path in paths:
-                qr_entry = tables.DatabaseEntry._from_query_result_block(block)
-                file_entries = list(
-                    tables.entries_from_file(path, self.default_waveunit))
-                for entry in file_entries:
-                    entry.source = qr_entry.source
-                    entry.provider = qr_entry.provider
-                    entry.physobs = qr_entry.physobs
-                    entry.fileid = qr_entry.fileid
-                    entry.observation_time_start =\
-                        qr_entry.observation_time_start
-                    entry.observation_time_end = qr_entry.observation_time_end
-                    entry.instrument = qr_entry.instrument
-                    entry.size = qr_entry.size
-                    entry.wavemin = qr_entry.wavemin
-                    entry.wavemax = qr_entry.wavemax
-                    entry.path = path
-                    entry.download_time = datetime.utcnow()
-                entries.extend(file_entries)
+        entries = list(self._download_and_collect_entries(
+            qr, client, path, progress))
         dump = serialize.dump_query(and_(*query))
         (dump_exists,), = self.session.query(
             exists().where(tables.JSONDump.dump == tables.JSONDump(dump).dump))
@@ -640,6 +646,30 @@ def add_from_hek_query_result(self, query_result,
             H2VClient().translate_and_query(query_result))
         self.add_from_vso_query_result(vso_qr, ignore_already_added)
 
+    def download_from_vso_query_result(self, query_result, client=None,
+            path=None, progress=False, ignore_already_added=False):
+        """download(query_result, client=sunpy.net.vso.VSOClient(),
+        path=None, progress=False, ignore_already_added=False)
+
+        Add new database entries from a VSO query result and download the
+        corresponding data files. See :meth:`sunpy.database.Database.download`
+        for information about the parameters `client`, `path`, `progress`.
+
+        Parameters
+        ----------
+        query_result : sunpy.net.vso.vso.QueryResponse
+            A VSO query response that was returned by the ``query`` method of a
+            :class:`sunpy.net.vso.VSOClient` object.
+
+        ignore_already_added : bool
+            See :meth:`sunpy.database.Database.add`.
+
+        """
+        if not query_result:
+            return
+        self.add_many(self._download_and_collect_entries(
+            query_result, client, path, progress))
+
     def add_from_vso_query_result(self, query_result,
             ignore_already_added=False):
         """Generate database entries from a VSO query result and add all the

sunpy/database/tests/test_database.py

@@ -49,6 +49,12 @@ def query_result():
         vso.attrs.Time('20130801T200000', '20130801T200030'),
         vso.attrs.Instrument('PLASTIC'))
 
+@pytest.fixture
+def download_qr():
+    return vso.VSOClient().query(
+        vso.attrs.Time('2012-03-29', '2012-03-29'),
+        vso.attrs.Instrument('AIA'))
+
 
 @pytest.fixture
 def empty_query():
@@ -401,6 +407,24 @@ def test_add_entry_from_hek_qr(database):
     database.add_from_hek_query_result(hek_res)
     assert len(database) == 2133
 
+
+@pytest.mark.online
+def test_download_from_qr(database, download_qr, tmpdir):
+    assert len(database) == 0
+    database.download_from_vso_query_result(
+        download_qr, path=str(tmpdir.join('{file}.fits')))
+    fits_pattern = str(tmpdir.join('*.fits'))
+    num_of_fits_headers = sum(
+        len(fits.get_header(file)) for file in glob.glob(fits_pattern))
+    assert len(database) == num_of_fits_headers > 0
+    for entry in database:
+        assert os.path.dirname(entry.path) == str(tmpdir)
+    database.undo()
+    assert len(database) == 0
+    database.redo()
+    assert len(database) == num_of_fits_headers > 0
+
+
 @pytest.mark.online
 def test_add_entry_from_qr(database, query_result):
     assert len(database) == 0