Merge pull request #3275 from zoghbi-a/fix-3271

BUG: fix Heasarch.locate_data returning empty rows

Author: bsipocz

Diff for: CHANGES.rst

@@ -18,6 +18,8 @@ heasarc
 ^^^^^^^
 
 - Add support for astropy.table.Row in Heasarc.download_data and Heasarc.locate_data. [#3270]
+- Heasarc.locate_data returns empty rows with an error in the error_message column if there are
+  no data associated with that row rather than filtering it out. [#3275]
 
 
 Infrastructure, Utility and Other Changes and Additions

Diff for: astroquery/heasarc/core.py

@@ -524,8 +524,13 @@ def locate_data(self, query_result=None, catalog_name=None):
             session=self._session
         )
         dl_result = query.execute().to_table()
-        dl_result = dl_result[dl_result['content_type'] == 'directory']
-        dl_result = dl_result[['ID', 'access_url', 'content_length']]
+        # include rows that have directory links (i.e. data) and those
+        # that report errors (usually means there are no data products)
+        dl_result = dl_result[np.ma.mask_or(
+            dl_result['content_type'] == 'directory',
+            dl_result['error_message'] != ''
+        )]
+        dl_result = dl_result[['ID', 'access_url', 'content_length', 'error_message']]
 
         # add sciserver and s3 columns
         newcol = [
@@ -622,6 +627,10 @@ def download_data(self, links, host='heasarc', location='.'):
                 '`~locate_data` first'
             )
 
+        # remove rows that dont have data, if any
+        if 'error_message' in links.colnames:
+            links = links[links['error_message'] == '']
+
         if host == 'heasarc':
 
             log.info('Downloading data from the HEASARC ...')

Diff for: astroquery/heasarc/tests/test_heasarc.py

@@ -365,6 +365,28 @@ def test_download_data__table_row():
         assert os.path.exists(f'{downloaddir}/data/file.txt')
 
 
+def test_download_data__exclude_rows_with_errors():
+    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'],
+            'error_message': ['', 'Error']
+        })
+        # 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, host="sciserver", location=downloaddir)
+        assert os.path.exists(f'{downloaddir}/file.txt')
+        # data/ should be excluded because it has an error
+        assert not os.path.exists(f'{downloaddir}/data')
+        assert not os.path.exists(f'{downloaddir}/data/file.txt')
+
+
 # S3 mock tests
 s3_bucket = "nasa-heasarc"
 s3_key1 = "some/location/file1.txt"