MAST: Adding read_product function to mast.observations

CHANGES.rst

@@ -174,7 +174,7 @@ mast
 - Added ``pass_id`` as an alias for the ``pass`` column in query functions for the Roman mission to avoid conflicts with
   the reserved Python keyword. [#3588]
 - Update the cutout format request parameter in ``Zcut.download_cutouts`` to reflect a recent service change. [#3608]
-
+- Adding in function to read FITS and ASDF data products to memory from STScI's S3 open data bucket using ``Observations.read_product()`` function. [#3610]
 
 jplspec
 ^^^^^^^

astroquery/mast/observations.py

@@ -11,9 +11,11 @@
 import warnings
 from pathlib import Path
 from urllib.parse import quote
+import importlib
 
 import astropy.coordinates as coord
 import astropy.units as u
+from astropy.io import fits
 import numpy as np
 from astropy.table import Row, Table, vstack
 from astropy.utils.decorators import deprecated_renamed_argument
@@ -41,6 +43,16 @@
 except ImportError:
     pass
 
+try:
+    import fsspec
+except ImportError:
+    fsspec = None
+
+try:
+    import asdf
+except ImportError:
+    asdf = None
+
 __all__ = ['Observations', 'ObservationsClass', 'MastClass', 'Mast']
 
 CLOUD_DISABLED_MESSAGE = (
@@ -1235,6 +1247,75 @@ def get_unique_product_list(self, observations, *, batch_size=500):
             log.info("To return all products, use `Observations.get_product_list`")
         return unique_products
 
+    # TODO: Need to inlcude way to parse if it is a MAST on prem URL and handle the streaming of that
+    def read_product(self, product_path, ignore_unrecognized=True, **kwargs):
+        """
+        Read a product from Open S3 bucket to memory. Currently supports FITS and ASDF product types only.
+
+        Parameters
+        ----------
+        product_path: str
+            URI to the product in the STScI S3 open data bucket.
+        ignore_unrecognized: bool
+            Tells asdf.open() to include or ignore warnings from unrecognized asdf tags. Defaults to True
+        **kwargs
+            Additional keyword arguments passed to the underlying file reader:
+            - For FITS files: forwarded to ``astropy.io.fits.open``.
+            Common options include ``memmap``, ``mode``, etc.
+            - Ignored for ASDF files (except for future extension if needed).
+
+        Returns
+        -------
+        object
+            FITS or ASDF object for the given data product.
+        """
+        # Checks if a path is empty or None.
+        if not product_path or not str(product_path).strip():
+            raise ValueError("No product path provided")
+
+        # Forces the path to be lowercase for the extension checks. This is only used for the checks
+        path = str(product_path).lower()
+
+        # Checks users environment for fsspec, required for both fits and asdf
+        if fsspec is None:
+            raise ImportError('The "fsspec" package is required to read products directly from a URI. '
+                              'Please install it with `pip install fsspec`.')
+
+        # Logic for reading FITS files
+        if path.endswith((".fits", ".fits.gz")):
+            try:
+                data_product = fits.open(product_path, fsspec_kwargs={"anon": True}, **kwargs)
+                log.info(f"Loaded: {product_path}")
+                return data_product
+            except Exception as e:
+                raise RuntimeError(f"Failed to open FITS File: {product_path} {e}")
+
+        # Logic for reading ASDF files
+        elif path.endswith(".asdf"):
+            if asdf is None:
+                raise ImportError('The "asdf" package is required to read ASDF files. Please install it with '
+                                  '`pip install asdf` or install astroquery with optional dependencies using '
+                                  '`pip install astroquery[all]`.')
+
+            for package in ["gwcs", "lz4", "roman_datamodels"]:
+                if importlib.util.find_spec(package) is None:
+                    warnings.warn(f'The "{package}" package is encouraged when reading ASDF files from the '
+                                  f'Roman Space Telescope mission. Please install it with `pip install {package}` '
+                                  f'or install astroquery with optional dependencies using '
+                                  f'`pip install astroquery[all]`.', ImportWarning)
+
+            # Attempts to open the asdf files
+            try:
+                f = fsspec.open(product_path, "rb", anon=True).open()
+                data_product = asdf.open(f, ignore_unrecognized_tag=ignore_unrecognized)
+                log.info(f"Loaded: {product_path}")
+                return data_product
+            except Exception as e:
+                raise RuntimeError(f"Failed to open ASDF File: {product_path} {e}")
+
+        else:
+            raise ValueError(f"Unsupported product type: {product_path}")
+
 
 @async_to_sync
 class MastClass(MastQueryWithLogin):

astroquery/mast/tests/test_mast.py

@@ -30,6 +30,12 @@
 except ImportError:
     pass
 
+try:
+    # Optional dependency import for ASDF file handling
+    import asdf
+except ImportError:
+    pass
+
 DATA_FILES = {'Mast.Caom.Cone': 'caom.json',
               'Mast.Name.Lookup': 'resolver.json',
               'mission_search_results': 'mission_results.json',
@@ -1373,6 +1379,102 @@ def test_observations_disable_cloud_dataset(patch_boto3):
     assert Observations._cloud_enabled_explicitly is False
 
 
+@pytest.fixture
+def mock_fits_open(mocker):
+    """Mock fits.open to return a valid HDUList without network access."""
+    return mocker.patch("astropy.io.fits.open", return_value=fits.HDUList([fits.PrimaryHDU()]))
+
+
+@pytest.fixture
+def mock_asdf_open(mocker):
+    pytest.importorskip("asdf")
+    pytest.importorskip("fsspec")
+
+    fake = mocker.Mock()
+    fake.open.return_value = "mock_asdf_file_object"
+    mocker.patch("fsspec.open", return_value=fake)
+
+    # Create a mock AsdfFile
+    if asdf is not None:
+        mock_asdf_file = MagicMock(spec=asdf.AsdfFile)
+    else:
+        mock_asdf_file = MagicMock()
+
+    return mocker.patch("asdf.open", return_value=mock_asdf_file)
+
+
+@pytest.mark.parametrize(
+    "product_path, expected_exception, match",
+    [
+        ("", ValueError, "No product path provided"),
+        ("   ", ValueError, "No product path provided"),
+        (None, ValueError, "No product path provided"),
+        ("unsupported_ex.txt", ValueError, "Unsupported product type"),
+    ],
+)
+def test_observations_read_product_invalid_inputs(product_path, expected_exception, match):
+    with pytest.raises(expected_exception, match=match):
+        Observations.read_product(product_path)
+
+
+@pytest.mark.parametrize(
+    ("module_name", "filename"),
+    [
+        ("fsspec", "file.fits"),
+        ("asdf", "file.asdf"),
+    ],
+)
+def test_observations_read_product_dependency_missing(monkeypatch, module_name, filename):
+    monkeypatch.setitem(Observations.read_product.__globals__, module_name, None)
+
+    with pytest.raises(ImportError, match=module_name):
+        Observations.read_product(filename)
+
+
+def test_observations_read_product_fits(mock_fits_open):
+    s3_fits_path = "s3://mock_fits_path.fits"
+    result = Observations.read_product(s3_fits_path)
+
+    mock_fits_open(s3_fits_path, fsspec_kwargs={"anon": True})
+    assert result is mock_fits_open.return_value
+
+
+def test_observations_read_product_asdf(mocker, mock_asdf_open):
+    # Mock importlib.util.find_spec to always return a spec (all packages present)
+    mocker.patch("importlib.util.find_spec", return_value=MagicMock())
+
+    s3_asdf_path = "s3://fake_asdf_path.asdf"
+    result = Observations.read_product(s3_asdf_path)
+
+    mock_asdf_open("mock_asdf_file_object")
+    assert result is mock_asdf_open.return_value
+
+
+def test_observations_read_product_asdf_missing_packages(mocker, mock_asdf_open):
+    """Test that warnings are issued for missing ASDF-related packages."""
+    # Mock importlib.util.find_spec to return None for the packages
+    def mock_find_spec(package_name):
+        if package_name in ["gwcs", "lz4", "roman_datamodels"]:
+            return None  # Simulate missing package
+        return MagicMock()  # Return something for other packages
+
+    mocker.patch("importlib.util.find_spec", side_effect=mock_find_spec)
+
+    # Capture warnings
+    with pytest.warns(ImportWarning) as warning_list:
+        obj = Observations.read_product("s3://fake_asdf_path.asdf")
+
+    # Check that warnings were issued for all three packages
+    assert len(warning_list) == 3
+    warning_messages = [str(w.message) for w in warning_list]
+
+    for package in ["gwcs", "lz4", "roman_datamodels"]:
+        assert any(package in msg and "encouraged" in msg for msg in warning_messages)
+
+    # Verify the object is still returned correctly
+    assert obj == mock_asdf_open.return_value
+
+
 ######################
 # CatalogClass tests #
 ######################

astroquery/mast/tests/test_mast_remote.py

@@ -18,6 +18,11 @@
 from ...exceptions import (InputWarning, InvalidQueryError, MaxResultsWarning, NoResultsWarning)
 from ..utils import ResolverError
 
+try:
+    import asdf
+except ImportError:
+    asdf = None
+
 
 @pytest.fixture(scope="module")
 def msa_product_table():
@@ -1048,6 +1053,23 @@ def test_observations_get_cloud_uris_no_duplicates(self, msa_product_table, rese
         uris = Observations.get_cloud_uris(products)
         assert len(uris) == 1
 
+    @pytest.mark.remote_data
+    @pytest.mark.parametrize(
+        "product_path, expected_type",
+        [
+            (
+                "s3://stpubdata/hst/public/u24r/u24r0102t/u24r0102t_c1f.fits",
+                fits.HDUList,
+            )
+        ],
+    )
+    def test_observations_read_product(self, product_path, expected_type):
+        pytest.importorskip("asdf")
+        pytest.importorskip("fsspec")
+
+        product = Observations.read_product(product_path)
+        assert isinstance(product, expected_type)
+
     ######################
     # CatalogClass tests #
     ######################

docs/conf.py

@@ -55,6 +55,8 @@
     'regions': ('https://astropy-regions.readthedocs.io/en/stable', None),
     'mocpy': ('https://cds-astro.github.io/mocpy', None),
     'pyvo': ('https://pyvo.readthedocs.io/en/stable', None),
+    'asdf': ('https://asdf.readthedocs.io/en/stable', None),
+    'fsspec': ('https://filesystem-spec.readthedocs.io/en/stable', None),
 })
 
 # -- Project information ------------------------------------------------------

docs/mast/mast_obsquery.rst

@@ -629,3 +629,47 @@ remain fully cloud-based.
    COMPLETE
    COMPLETE
    COMPLETE
+
+Streaming Data Products from S3 to memory
+-----------------------------------------
+If instead of downloading you would like to load an S3 URI directly to memory, you can use the `~astroquery.mast.ObservationsClass.read_product` method. 
+This function supports FITS and ASDF data products and will automatically parse the file for the suffix and load it to memory using `~astropy.io.fits.open` or `~asdf.open`.
+For ASDF data products, additional packages may be required such as ``roman_datamodels``, ``gwcs`` and ``lz4``.  To install astroquery with all these, and all other optional dependencies, use ``pip install astroquery[all]``.
+
+.. doctest-remote-data::
+
+   >>> from astroquery.mast import Observations
+   ...
+   >>> # For FITS files
+   >>> fits_product = Observations.read_product(product_path="s3://stpubdata/hst/public/u9o4/u9o40504m/u9o40504m_c3m.fits")
+   ... fits_product.info()
+   Filename: <class 's3fs.core.S3File'>
+   No.    Name      Ver    Type      Cards   Dimensions   Format
+   0  PRIMARY       1 PrimaryHDU       9   ()      
+   1                1 BinTableHDU     32   10000R x 5C   [1E, 1E, 1E, 1E, 1E]  
+   >>> # For ASDF files
+   >>> asdf_product = Observations.read_product(product_path="s3://stpubdata/roman/nexus/soc_simulations/tutorial_data/r0003201001001001004_0001_wfi01_f106_cal.asdf", ignore_unrecognized=True)
+   >>> asdf_product.info()
+   root (AsdfObject)
+   ├─asdf_library (Software)
+   │ ├─author (str): The ASDF Developers
+   │ ├─homepage (str): http://github.com/asdf-format/asdf
+   │ ├─name (str): asdf
+   │ └─version (str): 3.4.0
+   ├─history (dict)
+   │ └─extensions (list)
+   │   ├─[0] (ExtensionMetadata) ...
+   │   ├─[1] (ExtensionMetadata) ...
+   │   ├─[2] (ExtensionMetadata) ...
+   │   ├─[3] (ExtensionMetadata) ...
+   │   ├─[4] (ExtensionMetadata) ...
+   │   └─[5] (ExtensionMetadata) ...
+   └─roman (dict)
+   ├─catalogs (dict)
+   │ ├─galaxies (TaggedDict) ...
+   │ └─2 not shown
+   ├─data (NDArrayType) ...
+   ├─dq (NDArrayType) ...
+   ├─err (NDArrayType) ...
+   ├─meta (dict) ...
+   └─wcs (TaggedDict) ...

pyproject.toml

@@ -61,6 +61,7 @@ all = [
    "boto3",
    "botocore",
    "regions>=0.5",
+   "roman_datamodels",
 ]
 
 [build-system]

tox.ini

@@ -53,6 +53,9 @@ deps =
     oldestdeps-alldeps: mocpy==0.12
     oldestdeps-alldeps: regions==0.5
     oldestdeps-alldeps: astropy-healpix==0.7
+    oldestdeps-alldeps: roman_datamodels==0.11.0
+    oldestdeps-alldeps: gwcs==0.18
+
 
     online: pytest-custom_exit_code