S3 support in url loader/resolver (#3702)

Author: kecnry

jdaviz/configs/imviz/plugins/parsers.py

@@ -8,6 +8,7 @@
 from astropy.nddata import NDData
 from astropy.wcs import WCS
 from astropy.utils.data import cache_contents
+from urllib.parse import urlparse
 
 from glue.core.data import Component, Data
 from gwcs.wcs import WCS as GWCS
@@ -111,9 +112,9 @@ def parse_data(app, file_obj, ext=None, data_label=None,
         transformations.
     """
     if isinstance(file_obj, str):
-        file_obj = get_cloud_fits(
-            file_obj, ext=ext, cache=cache, local_path=local_path, timeout=timeout)
-        if not isinstance(file_obj, str):
+        if urlparse(file_obj).scheme == 's3':
+            file_obj = get_cloud_fits(
+                file_obj, ext=ext)
             _parse_image(
                 app, file_obj, data_label, ext=ext, parent=parent,
                 try_gwcs_to_fits_sip=gwcs_to_fits_sip

jdaviz/core/loaders/resolvers/url/url.py

@@ -7,7 +7,7 @@
 from jdaviz.core.registries import loader_resolver_registry
 from jdaviz.core.loaders.resolvers import BaseResolver
 from jdaviz.core.user_api import LoaderUserApi
-from jdaviz.utils import download_uri_to_path
+from jdaviz.utils import download_uri_to_path, get_cloud_fits
 
 
 __all__ = ['URLResolver']
@@ -19,6 +19,7 @@ class URLResolver(BaseResolver):
     default_input = 'url'
 
     url = Unicode("").tag(sync=True)
+    url_scheme = Unicode("").tag(sync=True)
     cache = Bool(True).tag(sync=True)
     local_path = Unicode("").tag(sync=True)
     timeout = FloatHandleEmpty(10).tag(sync=True)
@@ -33,7 +34,7 @@ def user_api(self):
 
     @property
     def is_valid(self):
-        return urlparse(self.url.strip()).scheme in ['http', 'https', 'mast', 'ftp']
+        return self.url_scheme in ['http', 'https', 'mast', 'ftp', 's3']
 
     @property
     def default_label(self):
@@ -45,6 +46,8 @@ def default_label(self):
 
     @observe('url', 'cache', 'timeout')
     def _on_url_changed(self, change):
+        self.url_scheme = urlparse(self.url.strip()).scheme
+
         # Clear the cached property to force re-download
         # or otherwise read from local file cache.
         if '_uri_output_file' in self.__dict__ and change['name'] in ('url', 'cache'):
@@ -54,6 +57,8 @@ def _on_url_changed(self, change):
 
     @cached_property
     def _uri_output_file(self):
+        if self.url_scheme == 's3':
+            return get_cloud_fits(self.url.strip())
         return download_uri_to_path(self.url.strip(), cache=self.cache,
                                     local_path=self.local_path, timeout=self.timeout)
 

jdaviz/core/loaders/resolvers/url/url.vue

@@ -23,7 +23,7 @@
       ></v-text-field>
     </v-row>
 
-    <v-row>
+    <v-row v-if="url_scheme !== 's3'">
       <v-text-field
         v-model.number='timeout'
         type="number"
@@ -35,6 +35,7 @@
     </v-row>
 
     <plugin-switch
+      v-if="url_scheme !== 's3'"
       :value.sync="cache"
       label="Cache File"
       api_hint="ldr.cache = "

jdaviz/core/loaders/test_loaders.py

@@ -121,6 +121,12 @@ def test_resolver_url(deconfigged_helper):
     loader.url = 'not-valid-url'
     assert len(loader.format.choices) == 0
 
+    # s3 input
+    loader.url = "s3://stpubdata/jwst/public/jw02727/L3/t/o002/jw02727-o002_t062_nircam_clear-f277w_i2d.fits"  # noqa: E501
+    assert loader._obj.url_scheme == 's3'
+    assert len(loader.format.choices) > 0
+
+    # https valid input
     loader.url = 'https://stsci.box.com/shared/static/exnkul627fcuhy5akf2gswytud5tazmw.fits'  # noqa
     assert len(loader.format.choices) == 4  # may change with future importers
     assert loader.format.selected == 'Image'  # default may change with future importers

jdaviz/utils.py

@@ -558,8 +558,7 @@ def __setgluestate__(cls, rec, context):
         return cls(masks=masks)
 
 
-def get_cloud_fits(possible_uri, ext=None, cache=None, local_path=os.curdir, timeout=None,
-                   dryrun=False):
+def get_cloud_fits(possible_uri, ext=None):
     """
     Retrieve and open a FITS file from an S3 URI using fsspec. Return the input
     unchanged if it is not an S3 URI.
@@ -579,10 +578,6 @@ def get_cloud_fits(possible_uri, ext=None, cache=None, local_path=os.curdir, tim
         Extension(s) to load from the FITS file. Can be an integer index (e.g., 0),
         a string name (e.g., "SCI"), or a list of such values. If `None`, all extensions
         are loaded.
-    cache : None, bool, or str, optional
-    local_path : str, optional
-    timeout : float, optional
-    dryrun : bool, optional
 
     Returns
     -------
@@ -594,24 +589,24 @@ def get_cloud_fits(possible_uri, ext=None, cache=None, local_path=os.curdir, tim
     parsed_uri = urlparse(possible_uri)
 
     # TODO: Add caching logic
-    if parsed_uri.scheme.lower() == 's3':
-        downloaded_hdus = []
-        # this loads the requested extensions into local memory:
-        with fits.open(possible_uri, fsspec_kwargs={"anon": True}) as hdul:
-            if ext is None:
-                ext_list = list(range(len(hdul)))
-            elif not isinstance(ext, list):
-                ext_list = [ext]
-            else:
-                ext_list = ext
-            for extension in ext_list:
-                hdu_obj = hdul[extension]
-                downloaded_hdus.append(hdu_obj.copy())
+    if not parsed_uri.scheme.lower() == 's3':
+        raise ValueError("Not an S3 URI: {}".format(possible_uri))
+
+    downloaded_hdus = []
+    # this loads the requested extensions into local memory:
+    with fits.open(possible_uri, fsspec_kwargs={"anon": True}) as hdul:
+        if ext is None:
+            ext_list = list(range(len(hdul)))
+        elif not isinstance(ext, list):
+            ext_list = [ext]
+        else:
+            ext_list = ext
+        for extension in ext_list:
+            hdu_obj = hdul[extension]
+            downloaded_hdus.append(hdu_obj.copy())
 
         file_obj = fits.HDUList(downloaded_hdus)
         return file_obj
-    # not s3 resource, return string as is
-    return possible_uri
 
 
 def cached_uri(uri):