diff --git a/CHANGES.rst b/CHANGES.rst index ae094673e4..113b7793d8 100644 --- a/CHANGES.rst +++ b/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 --git a/astroquery/heasarc/core.py b/astroquery/heasarc/core.py index 5827c936e9..a98b32873a 100644 --- a/astroquery/heasarc/core.py +++ b/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 --git a/astroquery/heasarc/tests/test_heasarc.py b/astroquery/heasarc/tests/test_heasarc.py index fc096b9682..a2f51b5c67 100644 --- a/astroquery/heasarc/tests/test_heasarc.py +++ b/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"