Merge pull request #3270 from zoghbi-a/heasarc-fix-3269

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

Author: bsipocz

CHANGES.rst

@@ -14,6 +14,10 @@ API changes
 Service fixes and enhancements
 ------------------------------
 
+heasarc
+^^^^^^^
+
+- Add support for astropy.table.Row in Heasarc.download_data and Heasarc.locate_data. [#3270]
 
 
 Infrastructure, Utility and Other Changes and Additions

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
@@ -470,7 +470,7 @@ def locate_data(self, query_result=None, catalog_name=None):
 
         Parameters
         ----------
-        query_result : `astropy.table.Table`, optional
+        query_result : `astropy.table.Table` or `astropy.table.Row`, optional
             A table that contain the search results. Typically as
             returned by query_region. If None, use the table from the
             most recent query_region call.
@@ -492,8 +492,14 @@ def locate_data(self, query_result=None, catalog_name=None):
             else:
                 query_result = self._last_result
 
-        if not isinstance(query_result, Table):
-            raise TypeError('query_result need to be an astropy.table.Table')
+        if isinstance(query_result, Row):
+            query_result = query_result.table[[query_result.index]]
+
+        elif not isinstance(query_result, Table):
+            raise TypeError(
+                'query_result need to be an astropy.table.Table or '
+                'astropy.table.Row'
+            )
 
         # make sure we have a column __row
         if '__row' not in query_result.colnames:
@@ -582,7 +588,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 +609,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

@@ -274,14 +274,28 @@ def test_locate_data():
         Heasarc.locate_data()
 
     with pytest.raises(
-        TypeError, match="query_result need to be an astropy.table.Table"
+        TypeError, match=(
+            "query_result need to be an astropy.table.Table or astropy.table.Row"
+        )
     ):
         Heasarc.locate_data([1, 2])
 
     with pytest.raises(ValueError, match="No __row column found"):
         Heasarc.locate_data(Table({"id": [1, 2, 3.0]}), catalog_name="xray")
 
 
+def test_locate_data_row():
+    table = Table({"id": [1, 2, 3.0]})
+
+    # we cannot do full call as this is not remote,
+    # but if we check that we error on __row not input type
+    with pytest.raises(ValueError, match="No __row column found"):
+        Heasarc.locate_data(table[0], catalog_name="xray")
+
+    with pytest.raises(ValueError, match="No __row column found"):
+        Heasarc.locate_data(table[0:2], catalog_name="xray")
+
+
 def test_download_data__empty():
     with pytest.raises(ValueError, match="Input links table is empty"):
         Heasarc.download_data(Table())
@@ -332,6 +346,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: