Merge pull request astropy#11860 from daria-cara/getdata-no-fallback

Remove fallback in getdata when ext specified

astropy/io/fits/convenience.py

@@ -136,6 +136,11 @@ def getdata(filename, *args, header=None, lower=None, upper=None, view=None,
 
             getdata('in.fits')
 
+        .. note::
+            Exclusive to ``getdata``: if extension is not specified
+            and primary header contains no data, ``getdata`` attempts
+            to retrieve data from first extension.
+
         By extension number::
 
             getdata('in.fits', 0)      # the primary header
@@ -186,22 +191,41 @@ def getdata(filename, *args, header=None, lower=None, upper=None, view=None,
 
         If the optional keyword ``header`` is set to `True`, this
         function will return a (``data``, ``header``) tuple.
+
+    Raises
+    ------
+    IndexError
+        If no data is found in searched extensions.
     """
 
     mode, closed = _get_file_mode(filename)
 
+    ext = kwargs.get('ext')
+    extname = kwargs.get('extname')
+    extver = kwargs.get('extver')
+    ext_given = not (len(args) == 0 and ext is None and
+                     extname is None and extver is None)
+
     hdulist, extidx = _getext(filename, mode, *args, **kwargs)
     try:
         hdu = hdulist[extidx]
         data = hdu.data
-        if data is None and extidx == 0:
-            try:
-                hdu = hdulist[1]
-                data = hdu.data
-            except IndexError:
-                raise IndexError('No data in this HDU.')
         if data is None:
-            raise IndexError('No data in this HDU.')
+            if ext_given:
+                raise IndexError(f"No data in HDU #{extidx}.")
+
+            # fallback to the first non-primary extension
+            if len(hdulist) == 1:
+                raise IndexError(
+                    "No data in Primary HDU and no extension HDU found."
+                    )
+            hdu = hdulist[1]
+            data = hdu.data
+            if data is None:
+                raise IndexError(
+                    "No data in either Primary or first extension HDUs."
+                    )
+
         if header:
             hdr = hdu.header
     finally:

astropy/io/fits/tests/test_convenience.py

@@ -3,6 +3,7 @@
 import os
 import pathlib
 import warnings
+import io
 
 import pytest
 import numpy as np
@@ -307,3 +308,85 @@ def test_pathlib(self):
 
         with fits.open(testfile) as hdul:
             np.testing.assert_array_equal(hdul[0].data, data)
+
+    def test_getdata_ext_given(self):
+        prihdu = fits.PrimaryHDU(data=np.zeros((5, 5), dtype=int))
+        exthdu1 = fits.ImageHDU(data=np.ones((5, 5), dtype=int))
+        exthdu2 = fits.ImageHDU(data=2 * np.ones((5, 5), dtype=int))
+        hdulist = fits.HDUList([prihdu, exthdu1, exthdu2])
+        buf = io.BytesIO()
+        hdulist.writeto(buf)
+
+        for ext in [0, 1, 2]:
+            buf.seek(0)
+            data = fits.getdata(buf, ext=ext)
+            assert data[0, 0] == ext
+
+    def test_getdata_ext_given_nodata(self):
+        prihdu = fits.PrimaryHDU(data=np.zeros((5, 5), dtype=int))
+        exthdu1 = fits.ImageHDU(data=np.ones((5, 5), dtype=int))
+        exthdu2 = fits.ImageHDU(data=None)
+        hdulist = fits.HDUList([prihdu, exthdu1, exthdu2])
+        buf = io.BytesIO()
+        hdulist.writeto(buf)
+        buf.seek(0)
+
+        with pytest.raises(IndexError, match="No data in HDU #2."):
+            fits.getdata(buf, ext=2)
+
+    def test_getdata_ext_not_given_with_data_in_primary(self):
+        prihdu = fits.PrimaryHDU(data=np.zeros((5, 5), dtype=int))
+        exthdu1 = fits.ImageHDU(data=None)
+        exthdu2 = fits.ImageHDU(data=None)
+        hdulist = fits.HDUList([prihdu, exthdu1, exthdu2])
+        buf = io.BytesIO()
+        hdulist.writeto(buf)
+        buf.seek(0)
+
+        data = fits.getdata(buf)
+        assert data[0, 0] == 0
+
+    def test_getdata_ext_not_given_with_data_in_ext(self):
+        # tests fallback mechanism
+        prihdu = fits.PrimaryHDU(data=None)
+        exthdu1 = fits.ImageHDU(data=np.ones((5, 5), dtype=int))
+        exthdu2 = fits.ImageHDU(data=None)
+        hdulist = fits.HDUList([prihdu, exthdu1, exthdu2])
+        buf = io.BytesIO()
+        hdulist.writeto(buf)
+        buf.seek(0)
+
+        data = fits.getdata(buf)
+        assert data[0, 0] == 1
+
+    def test_getdata_ext_not_given_nodata_any(self):
+        # tests exception raised when there is no data in either
+        # Primary HDU or first extension HDU
+        prihdu = fits.PrimaryHDU(data=None)
+        exthdu1 = fits.ImageHDU(data=None)
+        exthdu2 = fits.ImageHDU(data=np.ones((5, 5), dtype=int))
+        hdulist = fits.HDUList([prihdu, exthdu1, exthdu2])
+        buf = io.BytesIO()
+        hdulist.writeto(buf)
+        buf.seek(0)
+
+        with pytest.raises(
+            IndexError,
+            match="No data in either Primary or first extension HDUs."
+        ):
+            fits.getdata(buf)
+
+    def test_getdata_ext_not_given_nodata_noext(self):
+        # tests exception raised when there is no data in the
+        # Primary HDU and there are no extension HDUs
+        prihdu = fits.PrimaryHDU(data=None)
+        hdulist = fits.HDUList([prihdu])
+        buf = io.BytesIO()
+        hdulist.writeto(buf)
+        buf.seek(0)
+
+        with pytest.raises(
+            IndexError,
+            match="No data in Primary HDU and no extension HDU found."
+        ):
+            fits.getdata(buf)

docs/changes/io.fits/11860.bugfix.rst

@@ -0,0 +1,2 @@
+In ``fits.io.getdata`` do not fall back to first non-primary extension when
+user explicitly specifies an extension.