Return results in order of input, return as map

Author: snbianco

astroquery/mast/observations.py

@@ -874,7 +874,7 @@ def download_products(self, products, *, download_dir=None, flat=False,
         return manifest
 
     def get_cloud_uris(self, data_products=None, *, include_bucket=True, full_url=False, pagesize=None, page=None,
-                       mrp_only=False, extension=None, filter_products={}, **criteria):
+                       mrp_only=False, extension=None, filter_products={}, return_uri_map=False, **criteria):
         """
         Given an `~astropy.table.Table` of data products or query criteria and filter parameters,
         returns the associated cloud data URIs.
@@ -908,6 +908,10 @@ def get_cloud_uris(self, data_products=None, *, include_bucket=True, full_url=Fa
             or more acceptable values for that parameter.
             Filter behavior is AND between the filters and OR within a filter set.
             For example: {"productType": "SCIENCE", "extension"=["fits","jpg"]}
+        return_uri_map : bool, optional
+            Default False. If set to True, returns a dictionary mapping the original data product
+            URIs to their corresponding cloud URIs. This is useful for tracking which products were
+            successfully converted to cloud URIs.
         **criteria
             Criteria to apply. At least one non-positional criteria must be supplied.
             Valid criteria are coordinates, objectname, radius (as in `query_region` and `query_object`),
@@ -951,20 +955,31 @@ def get_cloud_uris(self, data_products=None, *, include_bucket=True, full_url=Fa
             # Filter product list
             data_products = self.filter_products(data_products, mrp_only=mrp_only, extension=extension,
                                                  **filter_products)
+            data_uris = data_products['dataURI']
         else:  # data_products is a list of URIs
             # Warn if trying to supply filters
             if filter_products or extension or mrp_only:
                 warnings.warn('Filtering is not supported when providing a list of MAST URIs. '
                               'To apply filters, please provide query criteria or a table of data products '
                               'as returned by `Observations.get_product_list`', InputWarning)
+            data_uris = data_products
 
-        if not len(data_products):
+        if not len(data_uris):
             warnings.warn('No matching products to fetch associated cloud URIs.', NoResultsWarning)
             return
 
         # Remove duplicate products
-        data_products = utils.remove_duplicate_products(data_products, 'dataURI')
-        return self._cloud_connection.get_cloud_uri_list(data_products, include_bucket, full_url)
+        data_uris = utils.remove_duplicate_products(data_uris, 'dataURI')
+
+        # Get cloud URIS
+        cloud_uris = self._cloud_connection.get_cloud_uri_list(data_uris, include_bucket, full_url)
+
+        # If return_uri_map is True, create a mapping of dataURIs to cloud URIs
+        if return_uri_map:
+            uri_map = dict(zip(data_uris, cloud_uris))
+            return uri_map
+
+        return cloud_uris
 
     def get_cloud_uri(self, data_product, *, include_bucket=True, full_url=False):
         """

astroquery/mast/tests/test_mast.py

@@ -353,7 +353,7 @@ def test_missions_get_product_list(patch_post):
 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()
+    assert (len(unique_products) == len(unique(unique_products, keys='filename')))
     # No INFO messages should be logged
     with caplog.at_level('INFO', logger='astroquery'):
         assert caplog.text == ''
@@ -770,6 +770,12 @@ def test_observations_get_cloud_uris(mock_client, patch_post):
     assert len(uris) == 1
     assert uris[0] == expected
 
+    # Return a map of URIs
+    uri_map = mast.Observations.get_cloud_uris([mast_uri], return_uri_map=True)
+    assert isinstance(uri_map, dict)
+    assert len(uri_map) == 1
+    assert uri_map[mast_uri] == expected
+
     # Warn if attempting to filter with list input
     with pytest.warns(InputWarning, match='Filtering is not supported'):
         mast.Observations.get_cloud_uris([mast_uri],

astroquery/mast/tests/test_mast_remote.py

@@ -269,7 +269,7 @@ def test_missions_get_unique_product_list(self, caplog):
         # Unique product list should have fewer rows
         assert len(products) > len(unique_products)
         # Rows should be unique based on filename
-        assert (unique_products == unique(unique_products, keys='filename')).all()
+        assert (len(unique_products) == len(unique(unique_products, keys='filename')))
         # Check that INFO messages were logged
         with caplog.at_level('INFO', logger='astroquery'):
             assert 'products were duplicates' in caplog.text
@@ -644,7 +644,7 @@ def test_observations_get_unique_product_list(self, caplog):
         # Unique product list should have fewer rows
         assert len(products) > len(unique_products)
         # Rows should be unique based on dataURI
-        assert (unique_products == unique(unique_products, keys='dataURI')).all()
+        assert (len(unique_products) == len(unique(unique_products, keys='dataURI')))
         # Check that INFO messages were logged
         with caplog.at_level('INFO', logger='astroquery'):
             assert 'products were duplicates' in caplog.text
@@ -878,6 +878,13 @@ def test_observations_get_cloud_uris_list_input(self):
         assert len(uris) > 0, f'Products for URI list {uri_list} were not found in the cloud.'
         assert uris == expected
 
+        # return map of dataURI to cloud URI
+        uri_map = Observations.get_cloud_uris(uri_list, return_uri_map=True)
+        assert isinstance(uri_map, dict)
+        assert len(uri_map) == 2
+        for i, uri in enumerate(uri_list):
+            assert uri_map[uri] == expected[i]
+
         # check for warning if filters are provided with list input
         with pytest.warns(InputWarning, match='Filtering is not supported'):
             Observations.get_cloud_uris(uri_list,

astroquery/mast/utils.py

@@ -13,7 +13,7 @@
 import platform
 
 from astropy.coordinates import SkyCoord
-from astropy.table import unique, Table
+from astropy.table import Table
 from astropy import units as u
 
 from .. import log
@@ -273,9 +273,9 @@ def mast_relative_path(mast_uri):
         The associated relative path(s).
     """
     if isinstance(mast_uri, str):
-        uri_list = [("uri", mast_uri)]
-    else:  # mast_uri parameter is a list
-        uri_list = [("uri", uri) for uri in mast_uri]
+        uri_list = [mast_uri]
+    else:
+        uri_list = list(mast_uri)
 
     # Split the list into chunks of 50 URIs; this is necessary
     # to avoid "414 Client Error: Request-URI Too Large".
@@ -284,19 +284,19 @@ def mast_relative_path(mast_uri):
     result = []
     for chunk in uri_list_chunks:
         response = _simple_request("https://mast.stsci.edu/api/v0.1/path_lookup/",
-                                   {"uri": [mast_uri[1] for mast_uri in chunk]})
+                                   {"uri": [mast_uri for mast_uri in chunk]})
 
         json_response = response.json()
 
         for uri in chunk:
             # Chunk is a list of tuples where the tuple is
             # ("uri", "/path/to/product")
             # so we index for path (index=1)
-            path = json_response.get(uri[1])["path"]
+            path = json_response.get(uri)["path"]
             if path is None:
-                warnings.warn(f"Failed to retrieve MAST relative path for {uri[1]}. Skipping...", NoResultsWarning)
-                continue
-            if 'galex' in path:
+                warnings.warn(f"Failed to retrieve MAST relative path for {uri}. Skipping...", NoResultsWarning)
+                path = None
+            elif 'galex' in path:
                 path = path.lstrip("/mast/")
             elif '/ps1/' in path:
                 path = path.replace("/ps1/", "panstarrs/ps1/public/")
@@ -330,20 +330,16 @@ def remove_duplicate_products(data_products, uri_key):
         Table containing products with unique dataURIs.
     """
     # Get unique products based on input type
+    seen = set()
     if isinstance(data_products, Table):
-        unique_products = unique(data_products, keys=uri_key)
-    else:  # data_products is a list
-        seen = set()
-        unique_products = []
-        for uri in data_products:
-            if uri not in seen:
-                seen.add(uri)
-                unique_products.append(uri)
-
-    number = len(data_products)
-    number_unique = len(unique_products)
-    if number_unique < number:
-        log.info(f"{number - number_unique} of {number} products were duplicates. "
-                 f"Only returning {number_unique} unique product(s).")
+        unique_rows = [row for row in data_products if not (row[uri_key] in seen or seen.add(row[uri_key]))]
+        unique_products = type(data_products)(rows=unique_rows)
+    else:  # Assume data_products is a list of URIs
+        unique_products = [uri for uri in data_products if not (uri in seen or seen.add(uri))]
+
+    duplicates_removed = len(data_products) - len(unique_products)
+    if duplicates_removed > 0:
+        log.info(f"{duplicates_removed} of {len(data_products)} products were duplicates. "
+                 f"Only returning {len(unique_products)} unique product(s).")
 
     return unique_products