fix-3269: add support for rows in heasarc.download_data

Author: zoghbi-a

astroquery/heasarc/core.py

@@ -6,7 +6,7 @@
 import tarfile
 import warnings
 import numpy as np
-from astropy.table import Table
+from astropy.table import Table, Row
 from astropy import coordinates
 from astropy import units as u
 from astropy.utils.decorators import deprecated, deprecated_renamed_argument
@@ -582,7 +582,7 @@ def download_data(self, links, host='heasarc', location='.'):
 
         Parameters
         ----------
-        links : `astropy.table.Table`
+        links : `astropy.table.Table` or `astropy.table.Row`
             The result from locate_data
         host : str
             The data host. The options are: heasarc (default), sciserver, aws.
@@ -603,6 +603,9 @@ def download_data(self, links, host='heasarc', location='.'):
         if len(links) == 0:
             raise ValueError('Input links table is empty')
 
+        if isinstance(links, Row):
+            links = links.table[[links.index]]
+
         if host not in ['heasarc', 'sciserver', 'aws']:
             raise ValueError('host has to be one of heasarc, sciserver, aws')
 

astroquery/heasarc/tests/test_heasarc.py

@@ -332,6 +332,25 @@ def test_download_data__outside_sciserver():
         )
 
 
+def test_download_data__table_row():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        datadir = f'{tmpdir}/data'
+        downloaddir = f'{tmpdir}/download'
+        os.makedirs(datadir, exist_ok=True)
+        with open(f'{datadir}/file.txt', 'w') as fp:
+            fp.write('data')
+        # include both a file and a directory
+        tab = Table({'sciserver': [f'{tmpdir}/data/file.txt', f'{tmpdir}/data']})
+        # The patch is to avoid the test that we are on sciserver
+        with patch('os.path.exists') as exists:
+            exists.return_value = True
+            Heasarc.download_data(tab[0], host="sciserver", location=downloaddir)
+            Heasarc.download_data(tab[1], host="sciserver", location=downloaddir)
+        assert os.path.exists(f'{downloaddir}/file.txt')
+        assert os.path.exists(f'{downloaddir}/data')
+        assert os.path.exists(f'{downloaddir}/data/file.txt')
+
+
 # S3 mock tests
 s3_bucket = "nasa-heasarc"
 s3_key1 = "some/location/file1.txt"

docs/heasarc/heasarc.rst

@@ -201,7 +201,7 @@ The first gives the url to the data from the main heasarc server. The second giv
 the local path to the data on Sciserver. The last gives the S3 URI to the data in the cloud.
 You can specify where the data are to be downloaded using the ``location`` parameter.
 
-To download the data, you can pass ``links`` table to `~astroquery.heasarc.HeasarcClass.download_data`,
+To download the data, you can pass ``links`` table (or row) to `~astroquery.heasarc.HeasarcClass.download_data`,
 specifying from where you want the data to be fetched by specifying the ``host`` parameter. By default,
 the data is fetched from the main HEASARC servers.
 The recommendation is to use different hosts depending on where your code is running: