More tests, download from HLSPs

snbianco · snbianco · commit eabbf774a7b7 · 2025-01-24T16:57:57.000-05:00
diff --git a/astroquery/mast/missions.py b/astroquery/mast/missions.py
@@ -42,12 +42,19 @@ class MastMissionsClass(MastQueryWithLogin):
     _search = 'search'
     _list_products = 'list_products'
 
+    # Workaround so that observation_id is returned in ULLYSES queries that do not specify columns
+    _default_ulysses_cols = ['target_name_ulysses', 'target_classification', 'targ_ra', 'targ_dec', 'host_galaxy_name',
+                             'spectral_type', 'bmv0_mag', 'u_mag', 'b_mag', 'v_mag', 'gaia_g_mean_mag', 'star_mass',
+                             'instrument', 'grating', 'filter', 'observation_id']
+
     def __init__(self, *, mission='hst'):
         super().__init__()
 
         self.dataset_kwds = {  # column keywords corresponding to dataset ID
             'hst': 'sci_data_set_name',
-            'jwst': 'fileSetName'
+            'jwst': 'fileSetName',
+            'classy': 'Target',
+            'ullyses': 'observation_id'
         }
 
         # Service attributes
@@ -69,7 +76,7 @@ def mission(self):
     @mission.setter
     def mission(self, value):
         # Need to update the service parameters if the mission is changed
-        self._mission = value
+        self._mission = value.lower()
         self._service_api_connection.set_service_params(self.service_dict, f'search/{self.mission}')
 
     def _parse_result(self, response, *, verbose=False):  # Used by the async_to_sync decorator functionality
@@ -184,6 +191,8 @@ def query_region_async(self, coordinates, *, radius=3*u.arcmin, limit=5000, offs
         # Dataset ID column should always be returned
         if select_cols:
             select_cols.append(self.dataset_kwds[self.mission])
+        elif self.mission == 'ullyses':
+            select_cols = self._default_ulysses_cols
 
         # basic params
         params = {'target': [f"{coordinates.ra.deg} {coordinates.dec.deg}"],
@@ -259,6 +268,8 @@ def query_criteria_async(self, *, coordinates=None, objectname=None, radius=3*u.
         # Dataset ID column should always be returned
         if select_cols:
             select_cols.append(self.dataset_kwds[self.mission])
+        elif self.mission == 'ullyses':
+            select_cols = self._default_ulysses_cols
 
         # build query
         params = {"limit": self.limit, "offset": offset, 'select_cols': select_cols}
@@ -452,10 +463,17 @@ def download_file(self, uri, *, local_path=None, cache=True, verbose=True):
             The full URL download path.
         """
 
-        # Construct the full data URL
-        base_url = self._service_api_connection.MISSIONS_DOWNLOAD_URL + self.mission + '/api/v0.1/retrieve_product'
-        data_url = base_url + '?product_name=' + uri
-        escaped_url = base_url + '?product_name=' + quote(uri, safe=':')
+        # Construct the full data URL based on mission
+        if self.mission in ['hst', 'jwst']:
+            # HST and JWST have a dedicated endpoint for retrieving products
+            base_url = self._service_api_connection.MISSIONS_DOWNLOAD_URL + self.mission + '/api/v0.1/retrieve_product'
+            keyword = 'product_name'
+        else:
+            # HLSPs use MAST download URL
+            base_url = self._service_api_connection.MAST_DOWNLOAD_URL
+            keyword = 'uri'
+        data_url = base_url + f'?{keyword}=' + uri
+        escaped_url = base_url + f'?{keyword}=' + quote(uri, safe='')
 
         # Determine local file path. Use current directory as default.
         filename = Path(uri).name
diff --git a/astroquery/mast/services.py b/astroquery/mast/services.py
@@ -14,6 +14,7 @@
 from astropy.table import Table, MaskedColumn
 from astropy.utils.decorators import deprecated_renamed_argument
 
+from .. import log
 from ..query import BaseQuery
 from ..utils import async_to_sync
 from ..utils.class_or_instance import class_or_instance
@@ -84,7 +85,12 @@ def _json_to_table(json_obj, data_key='data'):
             col_data = np.array([x[idx] for x in json_obj[data_key]], dtype=object)
         except KeyError:
             # it's not a data array, fall back to using column name as it is array of dictionaries
-            col_data = np.array([x[col_name] for x in json_obj[data_key]], dtype=object)
+            try:
+                col_data = np.array([x[col_name] for x in json_obj[data_key]], dtype=object)
+            except KeyError:
+                # Skip column names not found in data
+                log.debug('Column %s was not found in data. Skipping...', col_name)
+                continue
         if ignore_value is not None:
             col_data[np.where(np.equal(col_data, None))] = ignore_value
 
@@ -113,6 +119,7 @@ class ServiceAPI(BaseQuery):
     SERVICE_URL = conf.server
     REQUEST_URL = conf.server + "/api/v0.1/"
     MISSIONS_DOWNLOAD_URL = conf.server + "/search/"
+    MAST_DOWNLOAD_URL = conf.server + "/api/v0.1/Download/file"
     SERVICES = {}
 
     def __init__(self, session=None):
diff --git a/astroquery/mast/tests/data/README.rst b/astroquery/mast/tests/data/README.rst
@@ -25,3 +25,14 @@ To generate `~astroquery.mast.tests.data.panstarrs_columns.json`, use the follow
     >>> resp = utils._simple_request('https://catalogs.mast.stsci.edu/api/v0.1/panstarrs/dr2/mean/metadata.json')
     >>> with open('panstarrs_columns.json', 'w') as file:
     ...     json.dump(resp.json(), file, indent=4)  # doctest: +SKIP
+
+To generate `~astroquery.mast.tests.data.mission_products.json`, use the following:
+
+.. doctest-remote-data::
+
+    >>> import json
+    >>> from astroquery.mast import utils
+    ...
+    >>> resp = utils._simple_request('https://mast.stsci.edu/search/hst/api/v0.1/list_products', {'dataset_ids': 'Z14Z0104T'})
+    >>> with open('panstarrs_columns.json', 'w') as file:
+    ...     json.dump(resp.json(), file, indent=4)  # doctest: +SKIP
diff --git a/astroquery/mast/tests/data/mission_products.json b/astroquery/mast/tests/data/mission_products.json
@@ -0,0 +1,79 @@
+{
+    "products": [
+        {
+            "product_key": "Z14Z0104T_z14z0104t_shf.fits",
+            "access": "PUBLIC",
+            "dataset": "Z14Z0104T",
+            "instrument_name": "HRS   ",
+            "filters": "MIRROR-N2",
+            "filename": "z14z0104t_shf.fits",
+            "uri": "Z14Z0104T/z14z0104t_shf.fits",
+            "authz_primary_identifier": "Z14Z0104T",
+            "authz_secondary_identifier": "CAL",
+            "file_suffix": "SHF",
+            "category": "UNCALIBRATED",
+            "size": 31680,
+            "type": "science"
+        },
+        {
+            "product_key": "Z14Z0104T_z14z0104t_trl.fits",
+            "access": "PUBLIC",
+            "dataset": "Z14Z0104T",
+            "instrument_name": "HRS   ",
+            "filters": "MIRROR-N2",
+            "filename": "z14z0104t_trl.fits",
+            "uri": "Z14Z0104T/z14z0104t_trl.fits",
+            "authz_primary_identifier": "Z14Z0104T",
+            "authz_secondary_identifier": "CAL",
+            "file_suffix": "TRL",
+            "category": "AUX",
+            "size": 17280,
+            "type": "science"
+        },
+        {
+            "product_key": "Z14Z0104T_z14z0104t_ulf.fits",
+            "access": "PUBLIC",
+            "dataset": "Z14Z0104T",
+            "instrument_name": "HRS   ",
+            "filters": "MIRROR-N2",
+            "filename": "z14z0104t_ulf.fits",
+            "uri": "Z14Z0104T/z14z0104t_ulf.fits",
+            "authz_primary_identifier": "Z14Z0104T",
+            "authz_secondary_identifier": "CAL",
+            "file_suffix": "ULF",
+            "category": "UNCALIBRATED",
+            "size": 14400,
+            "type": "science"
+        },
+        {
+            "product_key": "Z14Z0104T_z14z0104t_pdq.fits",
+            "access": "PUBLIC",
+            "dataset": "Z14Z0104T",
+            "instrument_name": "HRS   ",
+            "filters": "MIRROR-N2",
+            "filename": "z14z0104t_pdq.fits",
+            "uri": "Z14Z0104T/z14z0104t_pdq.fits",
+            "authz_primary_identifier": "Z14Z0104T",
+            "authz_secondary_identifier": "PDQ",
+            "file_suffix": "PDQ",
+            "category": "AUX",
+            "size": 11520,
+            "type": "science"
+        },
+        {
+            "product_key": "Z14Z0104T_z14z0104x_ocx.fits",
+            "access": "PUBLIC",
+            "dataset": "Z14Z0104T",
+            "instrument_name": "HRS   ",
+            "filters": "MIRROR-N2",
+            "filename": "z14z0104x_ocx.fits",
+            "uri": "Z14Z0104T/z14z0104x_ocx.fits",
+            "authz_primary_identifier": "Z14Z0104X",
+            "authz_secondary_identifier": "OCX",
+            "file_suffix": "OCX",
+            "category": "OTHER",
+            "size": 11520,
+            "type": "science"
+        }
+    ]
+}
diff --git a/astroquery/mast/tests/test_mast.py b/astroquery/mast/tests/test_mast.py
@@ -7,22 +7,23 @@
 
 import pytest
 
-from astropy.table import Table
+from astropy.table import Table, unique
 from astropy.coordinates import SkyCoord
 from astropy.io import fits
 
 import astropy.units as u
 
 from astroquery.mast.services import _json_to_table
 from astroquery.utils.mocks import MockResponse
-from astroquery.exceptions import InvalidQueryError, InputWarning
+from astroquery.exceptions import InvalidQueryError, InputWarning, MaxResultsWarning, NoResultsWarning
 
 from astroquery import mast
 
 DATA_FILES = {'Mast.Caom.Cone': 'caom.json',
               'Mast.Name.Lookup': 'resolver.json',
               'mission_search_results': 'mission_results.json',
               'mission_columns': 'mission_columns.json',
+              'mission_products': 'mission_products.json',
               'columnsconfig': 'columnsconfig.json',
               'ticcolumns': 'ticcolumns.json',
               'ticcol_filtered': 'ticcolumns_filtered.json',
@@ -72,6 +73,7 @@ def patch_post(request):
     mp.setattr(mast.Observations, '_download_file', download_mockreturn)
     mp.setattr(mast.Observations, 'download_file', download_mockreturn)
     mp.setattr(mast.Catalogs, '_download_file', download_mockreturn)
+    mp.setattr(mast.MastMissions, '_download_file', download_mockreturn)
     mp.setattr(mast.Tesscut, '_download_file', tesscut_download_mockreturn)
     mp.setattr(mast.Zcut, '_download_file', zcut_download_mockreturn)
 
@@ -108,7 +110,7 @@ def post_mockreturn(self, method="POST", url=None, data=None, timeout=10, **kwar
     return [MockResponse(content)]
 
 
-def service_mockreturn(self, method="POST", url=None, data=None, timeout=10, use_json=False, **kwargs):
+def service_mockreturn(self, method="POST", url=None, data=None, params=None, timeout=10, use_json=False, **kwargs):
     if "panstarrs" in url:
         filename = data_path(DATA_FILES["panstarrs"])
     elif "tesscut" in url:
@@ -121,6 +123,8 @@ def service_mockreturn(self, method="POST", url=None, data=None, timeout=10, use
             filename = data_path(DATA_FILES['z_survey'])
         else:
             filename = data_path(DATA_FILES['z_cutout_fit'])
+    elif use_json and 'list_products' in url:
+        filename = data_path(DATA_FILES['mission_products'])
     elif use_json and data['radius'] == 300:
         filename = data_path(DATA_FILES["mission_incorrect_results"])
     elif use_json:
@@ -211,7 +215,9 @@ def test_missions_query_object(patch_post):
 
 
 def test_missions_query_region(patch_post):
-    result = mast.MastMissions.query_region(regionCoords, radius=0.002 * u.deg)
+    result = mast.MastMissions.query_region(regionCoords,
+                                            radius=0.002 * u.deg,
+                                            select_cols=['sci_pep_id'])
     assert isinstance(result, Table)
     assert len(result) > 0
 
@@ -242,6 +248,134 @@ def test_missions_query_criteria_async_with_missing_results(patch_post):
         _json_to_table(json.loads(responses), 'results')
 
 
+def test_missions_query_criteria(patch_post):
+    result = mast.MastMissions.query_criteria(
+        coordinates=regionCoords,
+        radius=3,
+        sci_pep_id=12556,
+        sci_obs_type='SPECTRUM',
+        sci_instrume='stis,acs,wfc3,cos,fos,foc,nicmos,ghrs',
+        sci_aec='S',
+        select_cols=['sci_pep_id', 'sci_instrume']
+    )
+    assert isinstance(result, Table)
+    assert len(result) > 0
+
+    # Raise error if non-positional criteria is not supplied
+    with pytest.raises(InvalidQueryError):
+        mast.MastMissions.query_criteria(
+            coordinates=regionCoords,
+            radius=3
+        )
+
+    # Raise error if invalid criteria is supplied
+    with pytest.raises(InvalidQueryError):
+        mast.MastMissions.query_criteria(
+            coordinates=regionCoords,
+            invalid=True
+        )
+
+    # Maximum results warning
+    with pytest.warns(MaxResultsWarning):
+        mast.MastMissions.query_criteria(
+            coordinates=regionCoords,
+            sci_aec='S',
+            limit=1
+        )
+
+
+def test_missions_get_product_list_async(patch_post):
+    # String input
+    result = mast.MastMissions.get_product_list_async('Z14Z0104T')
+    assert isinstance(result, MockResponse)
+
+    # List input
+    in_datasets = ['Z14Z0104T', 'Z14Z0102T']
+    result = mast.MastMissions.get_product_list_async(in_datasets)
+    assert isinstance(result, MockResponse)
+
+    # Row input
+    datasets = mast.MastMissions.query_object("M101", radius=".002 deg")
+    result = mast.MastMissions.get_product_list_async(datasets[:3])
+    assert isinstance(result, MockResponse)
+
+    # Table input
+    result = mast.MastMissions.get_product_list_async(datasets[0])
+    assert isinstance(result, MockResponse)
+
+    # Unsupported data type for datasets
+    with pytest.raises(TypeError) as err_type:
+        mast.MastMissions.get_product_list_async(1)
+    assert 'Unsupported data type' in str(err_type.value)
+
+    # Empty dataset list
+    with pytest.raises(InvalidQueryError) as err_empty:
+        mast.MastMissions.get_product_list_async([' '])
+    assert 'Dataset list is empty' in str(err_empty.value)
+
+
+def test_missions_get_product_list(patch_post):
+    # String input
+    result = mast.MastMissions.get_product_list('Z14Z0104T')
+    assert isinstance(result, Table)
+
+    # List input
+    in_datasets = ['Z14Z0104T', 'Z14Z0102T']
+    result = mast.MastMissions.get_product_list(in_datasets)
+    assert isinstance(result, Table)
+
+    # Row input
+    datasets = mast.MastMissions.query_object("M101", radius=".002 deg")
+    result = mast.MastMissions.get_product_list(datasets[:3])
+    assert isinstance(result, Table)
+
+    # Table input
+    result = mast.MastMissions.get_product_list(datasets[0])
+    assert isinstance(result, Table)
+
+
+def test_missions_get_unique_product_list(patch_post, caplog):
+    unique_products = mast.MastMissions.get_unique_product_list('Z14Z0104T')
+    assert isinstance(unique_products, Table)
+    assert (unique_products == unique(unique_products, keys='filename')).all()
+    # No INFO messages should be logged
+    with caplog.at_level('INFO', logger='astroquery'):
+        assert caplog.text == ''
+
+
+def test_missions_filter_products(patch_post):
+    # Filter products list by column
+    products = mast.MastMissions.get_product_list('Z14Z0104T')
+    filtered = mast.MastMissions.filter_products(products,
+                                                 category='CALIBRATED')
+    assert isinstance(filtered, Table)
+    assert all(filtered['category'] == 'CALIBRATED')
+
+    # Filter by non-existing column
+    with pytest.warns(InputWarning):
+        mast.MastMissions.filter_products(products,
+                                          invalid=True)
+
+
+def test_missions_download_products(patch_post, tmp_path):
+    # Check string input
+    test_dataset_id = 'Z14Z0104T'
+    result = mast.MastMissions.download_products(test_dataset_id,
+                                                 download_dir=tmp_path)
+    assert isinstance(result, Table)
+
+    # Check Row input
+    prods = mast.MastMissions.get_product_list('Z14Z0104T')
+    result = mast.MastMissions.download_products(prods[0],
+                                                 download_dir=tmp_path)
+    assert isinstance(result, Table)
+
+    # Warn about no products
+    with pytest.warns(NoResultsWarning):
+        result = mast.MastMissions.download_products(test_dataset_id,
+                                                     extension='jpg',
+                                                     download_dir=tmp_path)
+
 ###################
 # MastClass tests #
 ###################
diff --git a/astroquery/mast/tests/test_mast_remote.py b/astroquery/mast/tests/test_mast_remote.py
