diff --git a/CHANGES.rst b/CHANGES.rst index c6158ea041..c982493f2d 100644 --- a/CHANGES.rst +++ b/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 ^^^^^^^ diff --git a/astroquery/mast/observations.py b/astroquery/mast/observations.py index f516d40e20..844f8b3afe 100644 --- a/astroquery/mast/observations.py +++ b/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): diff --git a/astroquery/mast/tests/test_mast.py b/astroquery/mast/tests/test_mast.py index ad817949d3..d237c3a325 100644 --- a/astroquery/mast/tests/test_mast.py +++ b/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 # ###################### diff --git a/astroquery/mast/tests/test_mast_remote.py b/astroquery/mast/tests/test_mast_remote.py index a3c1d0405c..6ad5b746db 100644 --- a/astroquery/mast/tests/test_mast_remote.py +++ b/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 # ###################### diff --git a/docs/conf.py b/docs/conf.py index 037197122e..585a77a170 100644 --- a/docs/conf.py +++ b/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 ------------------------------------------------------ diff --git a/docs/mast/mast_obsquery.rst b/docs/mast/mast_obsquery.rst index 6ef2a65d47..4a7dc6a58a 100644 --- a/docs/mast/mast_obsquery.rst +++ b/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: + 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) ... diff --git a/pyproject.toml b/pyproject.toml index bbe6c2c6cd..6b1d794993 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,7 @@ all = [ "boto3", "botocore", "regions>=0.5", + "roman_datamodels", ] [build-system] diff --git a/tox.ini b/tox.ini index 1bd03f7e0d..217e1e8f37 100644 --- a/tox.ini +++ b/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