Merge pull request #2511 from jdavies-st/mast-no-subdirs

MAST: Add flat option to Observations.download_products

Author: bsipocz

CHANGES.rst

@@ -16,11 +16,6 @@ hsa
 
 - New module to access ESA Herschel mission. [#2122]
 
-mast
-^^^
-
-- Fixed ``Observations.get_product_list`` to handle input lists of obsids. [#2504]
-
 
 Service fixes and enhancements
 ------------------------------
@@ -95,6 +90,11 @@ mast
 - Cull duplicate downloads for the same dataURI in ``Observations.download_products()``
   and duplicate URIs in ``Observations.get_cloud_uris``. [#2497]
 
+- Fixed ``Observations.get_product_list`` to handle input lists of obsids. [#2504]
+
+- Add a ``flat`` option to ``Observation.download_products()`` to turn off the
+  automatic creation and organizing of products into subdirectories. [#2511]
+
 oac
 ^^^
 

astroquery/mast/observations.py

@@ -584,7 +584,7 @@ def download_file(self, uri, *, local_path=None, base_url=None, cache=True, clou
 
         return status, msg, url
 
-    def _download_files(self, products, base_dir, *, cache=True, cloud_only=False,):
+    def _download_files(self, products, base_dir, *, flat=False, cache=True, cloud_only=False,):
         """
         Takes an `~astropy.table.Table` of data products and downloads them into the directory given by base_dir.
 
@@ -594,6 +594,9 @@ def _download_files(self, products, base_dir, *, cache=True, cloud_only=False,):
             Table containing products to be downloaded.
         base_dir : str
             Directory in which files will be downloaded.
+        flat : bool
+            Default is False.  If set to True, no subdirectories will be made for the
+            downloaded files.
         cache : bool
             Default is True. If file is found on disk it will not be downloaded again.
         cloud_only : bool, optional
@@ -610,9 +613,12 @@ def _download_files(self, products, base_dir, *, cache=True, cloud_only=False,):
         for data_product in products:
 
             # create the local file download path
-            local_path = os.path.join(base_dir, data_product['obs_collection'], data_product['obs_id'])
-            if not os.path.exists(local_path):
-                os.makedirs(local_path)
+            if not flat:
+                local_path = os.path.join(base_dir, data_product['obs_collection'], data_product['obs_id'])
+                if not os.path.exists(local_path):
+                    os.makedirs(local_path)
+            else:
+                local_path = base_dir
             local_path = os.path.join(local_path, os.path.basename(data_product['productFilename']))
 
             # download the files
@@ -642,8 +648,8 @@ def _download_curl_script(self, products, out_dir):
         """
 
         url_list = [("uri", url) for url in products['dataURI']]
-        download_file = "mastDownload_" + time.strftime("%Y%m%d%H%M%S")
-        local_path = os.path.join(out_dir.rstrip('/'), download_file + ".sh")
+        download_file = "mastDownload_" + time.strftime("%Y%m%d%H%M%S") + ".sh"
+        local_path = os.path.join(out_dir, download_file)
 
         response = self._download_file(self._portal_api_connection.MAST_BUNDLE_URL + ".sh",
                                        local_path, data=url_list, method="POST")
@@ -660,7 +666,7 @@ def _download_curl_script(self, products, out_dir):
                           'Message': [msg]})
         return manifest
 
-    def download_products(self, products, *, download_dir=None,
+    def download_products(self, products, *, download_dir=None, flat=False,
                           cache=True, curl_flag=False, mrp_only=False, cloud_only=False, **filters):
         """
         Download data products.
@@ -673,6 +679,14 @@ def download_products(self, products, *, download_dir=None,
             or a Table of products (as is returned by `get_product_list`)
         download_dir : str, optional
             Optional.  Directory to download files to.  Defaults to current directory.
+        flat : bool, optional
+            Default is False.  If set to True, and download_dir is specified, it will put
+            all files into download_dir without subdirectories.  Or if set to True and
+            download_dir is not specified, it will put files in the current directory,
+            again with no subdirs.  The default of False puts files into the standard
+            directory structure of "mastDownload/<obs_collection>/<obs_id>/".  If
+            curl_flag=True, the flat flag has no effect, as astroquery does not control
+            how MAST generates the curl download script.
         cache : bool, optional
             Default is True. If file is found on disc it will not be downloaded again.
             Note: has no affect when downloading curl script.
@@ -731,13 +745,19 @@ def download_products(self, products, *, download_dir=None,
             download_dir = '.'
 
         if curl_flag:  # don't want to download the files now, just the curl script
+            if flat:
+                # flat=True doesn't work with curl_flag=True, so issue a warning
+                warnings.warn("flat=True has no effect on curl downloads.", InputWarning)
             manifest = self._download_curl_script(products,
                                                   download_dir)
 
         else:
-            base_dir = download_dir.rstrip('/') + "/mastDownload"
+            if flat:
+                base_dir = download_dir
+            else:
+                base_dir = os.path.join(download_dir, "mastDownload")
             manifest = self._download_files(products,
-                                            base_dir=base_dir,
+                                            base_dir=base_dir, flat=flat,
                                             cache=cache,
                                             cloud_only=cloud_only)
 

astroquery/mast/tests/test_mast_remote.py

@@ -1,5 +1,7 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 
+from re import sub
+from pathlib import Path
 import numpy as np
 import os
 import pytest
@@ -14,13 +16,26 @@
 from astroquery import mast
 
 from ..utils import ResolverError
-from ...exceptions import (InvalidQueryError, MaxResultsWarning, NoResultsWarning,
+from ...exceptions import (InputWarning, InvalidQueryError, MaxResultsWarning, NoResultsWarning,
                            RemoteServiceError)
 
 
 OBSID = '1647157'
 
 
+@pytest.fixture(scope="module")
+def msa_product_table():
+    # Pull products for a JWST NIRSpec MSA observation with 6 known
+    # duplicates of the MSA configuration file, propID=2736
+    products = mast.Observations.get_product_list("87602009")
+
+    # Filter out everything but the MSA config file
+    mask = np.char.find(products["dataURI"], "_msa.fits") != -1
+    products = products[mask]
+
+    return products
+
+
 @pytest.mark.remote_data
 class TestMast:
 
@@ -290,21 +305,42 @@ def test_observations_download_products(self, tmpdir):
         assert os.path.isfile(result2['Local Path'][0])
         assert len(result2) == 1
 
-    def test_observations_download_products_no_duplicates(self, tmpdir, caplog):
+    def test_observations_download_products_flat(self, tmp_path, msa_product_table):
+
+        # Get a product list with 6 duplicate JWST MSA config files
+        products = msa_product_table
+
+        assert len(products) == 6
+
+        # Download with flat=True
+        manifest = mast.Observations.download_products(products, flat=True,
+                                                       download_dir=tmp_path)
+
+        assert Path(manifest["Local Path"][0]).parent == tmp_path
+
+    def test_observations_download_products_flat_curl(self, tmp_path, msa_product_table):
+
+        # Get a product list with 6 duplicate JWST MSA config files
+        products = msa_product_table
 
-        # Pull products for a JWST NIRSpec MSA observation with 6 known
-        # duplicates of the MSA configuration file, propID=2736
-        products = mast.Observations.get_product_list("87602009")
+        assert len(products) == 6
+
+        # Download with flat=True, curl_flag=True, look for warning
+        with pytest.warns(InputWarning):
+            mast.Observations.download_products(products, flat=True,
+                                                curl_flag=True,
+                                                download_dir=tmp_path)
+
+    def test_observations_download_products_no_duplicates(self, tmp_path, caplog, msa_product_table):
 
-        # Filter out everything but the MSA config file
-        mask = np.char.find(products["dataURI"], "_msa.fits") != -1
-        products = products[mask]
+        # Get a product list with 6 duplicate JWST MSA config files
+        products = msa_product_table
 
         assert len(products) == 6
 
         # Download the product
         manifest = mast.Observations.download_products(products,
-                                                       download_dir=str(tmpdir))
+                                                       download_dir=tmp_path)
 
         # Check that it downloads the MSA config file only once
         assert len(manifest) == 1
@@ -313,14 +349,24 @@ def test_observations_download_products_no_duplicates(self, tmpdir, caplog):
         with caplog.at_level("INFO", logger="astroquery"):
             assert "products were duplicates" in caplog.text
 
+    def test_observations_get_cloud_uris_no_duplicates(self, msa_product_table):
+
+        # Get a product list with 6 duplicate JWST MSA config files
+        products = msa_product_table
+
+        assert len(products) == 6
+
         # enable access to public AWS S3 bucket
-        mast.Observations.enable_cloud_dataset()
+        mast.Observations.enable_cloud_dataset(provider='AWS')
 
-        # Check duplicate cloud URIs as well
-        uris = mast.Observations.get_cloud_uris(products)
+        # Check for cloud URIs.  Accept a NoResultsWarning if AWS S3
+        # doesn't have the file.  It doesn't matter as we're only checking
+        # that the duplicate products have been culled to a single one.
+        with pytest.warns(NoResultsWarning):
+            uris = mast.Observations.get_cloud_uris(products)
         assert len(uris) == 1
 
-    def test_observations_download_file(self, tmpdir):
+    def test_observations_download_file(self):
 
         # enabling cloud connection
         mast.Observations.enable_cloud_dataset(provider='AWS')