check TBIT for slices/rows read

esheldon · Aug 5, 2024 · 46b7a57 · 46b7a57
1 parent 5b2cb45
commit 46b7a57
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 0 deletions.
diff --git a/fitsio/hdu/table.py b/fitsio/hdu/table.py
@@ -896,6 +896,8 @@ def read_rows(self, rows, vstorage=None,
         dtype, offsets, isvar = self.get_rec_dtype(vstorage=vstorage)
 
         w, = np.where(isvar == True)  # noqa
+        has_tbit = self._check_tbit()
+
         if w.size > 0:
             if vstorage is None:
                 _vstorage = self._vstorage
@@ -905,6 +907,14 @@ def read_rows(self, rows, vstorage=None,
             return self._read_rec_with_var(
                 colnums, rows, sortind, dtype, offsets, isvar, _vstorage,
             )
+        elif has_tbit:
+            # drop down to read_columns since we can't stuff into a
+            # contiguous array
+            colnums = self._extract_colnums()
+            array = self.read_columns(
+                colnums, rows=rows, vstorage=vstorage, upper=upper,
+                lower=lower, trim_strings=trim_strings,
+            )
         else:
             array = np.zeros(rows.size, dtype=dtype)
             self._FITS.read_rows_as_rec(self._ext+1, array, rows, sortind)
@@ -1094,6 +1104,8 @@ def read_slice(self, firstrow, lastrow, step=1,
         dtype, offsets, isvar = self.get_rec_dtype(vstorage=vstorage)
 
         w, = np.where(isvar == True)  # noqa
+        has_tbit = self._check_tbit()
+
         if w.size > 0:
             if vstorage is None:
                 _vstorage = self._vstorage
@@ -1104,6 +1116,15 @@ def read_slice(self, firstrow, lastrow, step=1,
             colnums = self._extract_colnums()
             array = self._read_rec_with_var(
                 colnums, rows, sortind, dtype, offsets, isvar, _vstorage)
+        elif has_tbit:
+            # drop down to read_columns since we can't stuff into a
+            # contiguous array
+            colnums = self._extract_colnums()
+            rows = np.arange(firstrow, lastrow, step, dtype='i8')
+            array = self.read_columns(
+                colnums, rows=rows, vstorage=vstorage, upper=upper,
+                lower=lower, trim_strings=trim_strings,
+            )
         else:
             if step != 1:
                 rows = np.arange(firstrow, lastrow, step, dtype='i8')

diff --git a/fitsio/tests/test_table.py b/fitsio/tests/test_table.py
@@ -1257,6 +1257,13 @@ def test_table_bitcol_read_write():
             d = fits[1].read()
             compare_rec(bdata, d, "table read/write")
 
+            rows = [0, 2]
+            d = fits[1].read(rows=rows)
+            compare_rec(bdata[rows], d, "table read/write rows")
+
+            d = fits[1][:2]
+            compare_rec(bdata[:2], d, "table read/write slice")
+
         # now test read_column
         with FITS(fname) as fits: