MAST: Adding read_product function to mast.observations

CHANGES.rst

@@ -173,7 +173,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

@@ -14,6 +14,7 @@
 
 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 +42,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 +1246,72 @@ 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 enviornment 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 fsspec is None:
+                raise ImportError('The "fsspec" package is required to read products directly from a URI. '
+                                  'Please install it with `pip install fsspec` or install astroquery with '
+                                  'optional dependencies using `pip install astroquery[all]`.')
+            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]`.')
+
+            # 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

@@ -1372,6 +1372,72 @@ 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):
+    return mocker.patch(
+        "asdf.open",
+        return_value=MagicMock(name="AsdfFile"),
+    )
+
+
+@pytest.fixture
+def mock_fsspec_open(mocker):
+    fake = mocker.Mock()
+    fake.open.return_value = "mock_asdf_file_object"
+    return mocker.patch("fsspec.open", return_value=fake)
+
+
+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(mock_asdf_open, mock_fsspec_open):
+    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
+
+
+@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)
+
+
+def test_observations_read_product_fsspec_missing(monkeypatch):
+    # Forces fsspec to be None
+    monkeypatch.setitem(Observations.read_product.__globals__, "fsspec", None)
+
+    with pytest.raises(ImportError, match="fsspec"):
+        Observations.read_product("file.fits")
+
+
+def test_observations_read_product_asdf_missing(monkeypatch):
+    # Forces asdf to be None
+    monkeypatch.setitem(Observations.read_product.__globals__, "asdf", None)
+
+    with pytest.raises(ImportError, match="asdf"):
+        Observations.read_product("file.asdf")
+
 ######################
 # CatalogClass tests #
 ######################

astroquery/mast/tests/test_mast_remote.py

@@ -20,6 +20,11 @@
 from ...exceptions import (InputWarning, InvalidQueryError, MaxResultsWarning,
                            NoResultsWarning)
 
+try:
+    import asdf
+except ImportError:
+    asdf = None
+
 
 @pytest.fixture(scope="module")
 def msa_product_table():
@@ -1048,6 +1053,25 @@ 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,
+            ),
+            (
+                "s3://stpubdata/roman/nexus/open_universe/WAS/simple_model/"
+                "F184/10073/roman_was_F184_10073_wfi02_simple.asdf",
+                asdf.AsdfFile,
+            ),
+        ],
+    )
+    def test_observations_read_product(self, product_path, expected_type):
+        product = Observations.read_product(product_path)
+        assert isinstance(product, expected_type)
+
     ######################
     # CatalogClass tests #
     ######################

docs/conf.py

@@ -56,6 +56,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,15 @@ 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
+   >>> fits_product = Observations.read_product(product_path="s3://stpubdata/hst/public/u9o4/u9o40504m/u9o40504m_c3m.fits")
+   >>> asdf_product = Observations.read_product(product_path="s3://stpubdata/roman/nexus/soc_simulations/tutorial_data/r0003201001001001004_0001_wfi01_f106_cal.asdf", ignore_unrecognized=True)

pyproject.toml

@@ -42,9 +42,9 @@ test = [
    # Temp workaround for https://github.com/RKrahl/pytest-dependency/issues/91
    "pytest-dependency; platform_system != 'Windows'",
    "pytest-rerunfailures",
-   "fsspec[http]",
+   "fsspec[http,s3]",
    "moto[s3]",
-   "s3fs",
+   "asdf",
 ]
 docs = [
    "sphinx",
@@ -54,13 +54,19 @@ docs = [
    "sphinx-astropy>=1.5",
    "scipy",
    "fsspec[http]",
+   "asdf"
 ]
 all = [
    "mocpy>=0.12",
    "astropy-healpix",
    "boto3",
    "botocore",
    "regions>=0.5",
+   "asdf",
+   "lz4",
+   "gwcs",
+   "roman_datamodels",
+   "fsspec[http,s3]"
 ]
 
 [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