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..a84e806e5f 100644 --- a/astroquery/mast/observations.py +++ b/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): diff --git a/astroquery/mast/tests/test_mast.py b/astroquery/mast/tests/test_mast.py index ad817949d3..6365d4844c 100644 --- a/astroquery/mast/tests/test_mast.py +++ b/astroquery/mast/tests/test_mast.py @@ -1373,6 +1373,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 # ###################### diff --git a/astroquery/mast/tests/test_mast_remote.py b/astroquery/mast/tests/test_mast_remote.py index b6f55a05e2..a617a39228 100644 --- a/astroquery/mast/tests/test_mast_remote.py +++ b/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 # ###################### diff --git a/docs/conf.py b/docs/conf.py index 2178e00ff4..623c302323 100644 --- a/docs/conf.py +++ b/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 ------------------------------------------------------ diff --git a/docs/mast/mast_obsquery.rst b/docs/mast/mast_obsquery.rst index 6ef2a65d47..3ad0e70728 100644 --- a/docs/mast/mast_obsquery.rst +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index bbe6c2c6cd..e789f52e81 100644 --- a/pyproject.toml +++ b/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,6 +54,7 @@ docs = [ "sphinx-astropy>=1.5", "scipy", "fsspec[http]", + "asdf" ] all = [ "mocpy>=0.12", @@ -61,6 +62,11 @@ all = [ "boto3", "botocore", "regions>=0.5", + "asdf", + "lz4", + "gwcs", + "roman_datamodels", + "fsspec[http,s3]" ] [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