Merge pull request #14428 from astrofrog/tile-size-property

Expose tile shape and compression type on CompImageHDU

astropy/io/fits/_tiled_compression/tests/test_fitsio.py

@@ -59,30 +59,30 @@
     ids=lambda x: f"shape: {x[0]} tile_dims: {x[1]}",
 )
 def array_shapes_tile_dims(request, compression_type):
-    shape, tile_dim = request.param
+    shape, tile_dims = request.param
     # H_COMPRESS needs >=2D data and always 2D tiles
     if compression_type == "HCOMPRESS_1":
         if (
             # We don't have at least a 2D image
             len(shape) < 2
             or
             # We don't have 2D tiles
-            np.count_nonzero(np.array(tile_dim) != 1) != 2
+            np.count_nonzero(np.array(tile_dims) != 1) != 2
             or
             # TODO: The following restrictions can be lifted with some extra work.
             # The tile is not the first two dimensions of the data
-            tile_dim[0] == 1
-            or tile_dim[1] == 1
+            tile_dims[0] == 1
+            or tile_dims[1] == 1
             or
             # The tile dimensions not an integer multiple of the array dims
-            np.count_nonzero(np.array(shape[:2]) % tile_dim[:2]) != 0
+            np.count_nonzero(np.array(shape[:2]) % tile_dims[:2]) != 0
         ):
             pytest.xfail(
                 "HCOMPRESS requires 2D tiles, from the first two"
                 "dimensions, and an integer number of tiles along the first two"
                 "axes."
             )
-    return shape, tile_dim
+    return shape, tile_dims
 
 
 @pytest.fixture(scope="module")
@@ -170,8 +170,7 @@ def astropy_compressed_file_path(
     hdu = fits.CompImageHDU(
         data=original_data,
         compression_type=compression_type,
-        # TODO: why does this require a list??
-        tile_size=list(tile_dims) if tile_dims is not None else tile_dims,
+        tile_shape=None if tile_dims is None else tile_dims[::-1],
         **param,
     )
     hdu.writeto(filename)

astropy/io/fits/_tiled_compression/tests/test_tiled_compression.py

@@ -52,7 +52,9 @@ def test_zblank_support(canonical_data_base_path, tmp_path):
     # Now generate a file ourselves and check that the output has the ZBLANK
     # keyword set automatically
 
-    hdu = fits.CompImageHDU(data=reference, compression_type="RICE_1", tile_size=(6, 6))
+    hdu = fits.CompImageHDU(
+        data=reference, compression_type="RICE_1", tile_shape=(6, 6)
+    )
 
     hdu.writeto(tmp_path / "test_zblank.fits")
 
@@ -62,7 +64,7 @@ def test_zblank_support(canonical_data_base_path, tmp_path):
 
 
 @pytest.mark.parametrize(
-    ("shape", "tile_dim"),
+    ("shape", "tile_shape"),
     (
         ([10, 10], [5, 5]),  # something for HCOMPRESS
         ([5, 5, 5], [5, 5, 5]),
@@ -75,22 +77,22 @@ def test_zblank_support(canonical_data_base_path, tmp_path):
     ),
 )
 def test_roundtrip_high_D(
-    numpy_rng, compression_type, compression_param, tmp_path, dtype, shape, tile_dim
+    numpy_rng, compression_type, compression_param, tmp_path, dtype, shape, tile_shape
 ):
     if compression_type == "HCOMPRESS_1" and (
         # We don't have at least a 2D image
         len(shape) < 2
         or
         # We don't have 2D tiles
-        np.count_nonzero(np.array(tile_dim) != 1) != 2
+        np.count_nonzero(np.array(tile_shape) != 1) != 2
         or
         # TODO: The following restrictions can be lifted with some extra work.
         # The tile is not the first two dimensions of the data
-        tile_dim[0] == 1
-        or tile_dim[1] == 1
+        tile_shape[0] == 1
+        or tile_shape[1] == 1
         or
         # The tile dimensions not an integer multiple of the array dims
-        np.count_nonzero(np.array(shape[:2]) % tile_dim[:2]) != 0
+        np.count_nonzero(np.array(shape[:2]) % tile_shape[:2]) != 0
     ):
         pytest.xfail("HCOMPRESS requires 2D tiles.")
     random = numpy_rng.uniform(high=255, size=shape)
@@ -112,7 +114,7 @@ def test_roundtrip_high_D(
     hdu = fits.CompImageHDU(
         data=original_data,
         compression_type=compression_type,
-        tile_size=tile_dim,
+        tile_shape=tile_shape,
         **param,
     )
     hdu.writeto(filename)

astropy/io/fits/_tiled_compression/tiled_compression.py

@@ -303,7 +303,7 @@ def decompress_hdu_section(hdu, first_tile_index, last_tile_index):
             )
 
         else:
-            if hdu._header["ZCMPTYPE"] == "GZIP_2":
+            if hdu.compression_type == "GZIP_2":
                 # Decompress with GZIP_1 just to find the total number of
                 # elements in the uncompressed data.
                 # TODO: find a way to avoid doing this for all tiles
@@ -313,14 +313,14 @@ def decompress_hdu_section(hdu, first_tile_index, last_tile_index):
                 )
 
             tile_buffer = _decompress_tile(
-                cdata, algorithm=hdu._header["ZCMPTYPE"], **settings
+                cdata, algorithm=hdu.compression_type, **settings
             )
 
             tile_data = _finalize_array(
                 tile_buffer,
                 bitpix=hdu._header["ZBITPIX"],
                 tile_shape=actual_tile_shape,
-                algorithm=hdu._header["ZCMPTYPE"],
+                algorithm=hdu.compression_type,
                 lossless=not quantized,
             )
 
@@ -464,7 +464,7 @@ def compress_hdu(hdu):
         if gzip_fallback[-1]:
             cbytes = _compress_tile(data, algorithm="GZIP_1")
         else:
-            cbytes = _compress_tile(data, algorithm=hdu._header["ZCMPTYPE"], **settings)
+            cbytes = _compress_tile(data, algorithm=hdu.compression_type, **settings)
         compressed_bytes.append(cbytes)
 
     if zblank is not None:
@@ -489,11 +489,11 @@ def compress_hdu(hdu):
 
     # For PLIO_1, the size of each heap element is a factor of two lower than
     # the real size - not clear if this is deliberate or bug somewhere.
-    if hdu._header["ZCMPTYPE"] == "PLIO_1":
+    if hdu.compression_type == "PLIO_1":
         table["COMPRESSED_DATA"][:, 0] //= 2
 
     # For PLIO_1, it looks like the compressed data is always stored big endian
-    if hdu._header["ZCMPTYPE"] == "PLIO_1":
+    if hdu.compression_type == "PLIO_1":
         for irow in range(len(compressed_bytes)):
             if not gzip_fallback[irow]:
                 array = np.frombuffer(compressed_bytes[irow], dtype="i2")