Properly set color interpretation after writing a colormap (#3136)

* Properly set color interpretation after writing a colormap

Resolves #3133

* Appease flake8

CHANGES.txt

@@ -1,6 +1,14 @@
 Changes
 =======
 
+Next (TBD)
+----------
+
+Bug fixes:
+
+- Color interpretation is correctly set to "palette" after a colormap is
+  written (#3133).
+
 1.4b1 (2024-08-10)
 ------------------
 

rasterio/_io.pyx

@@ -2036,7 +2036,7 @@ cdef class DatasetWriterBase(DatasetReaderBase):
                 GDALSetColorEntry(hTable, i, &color)
 
             # TODO: other color interpretations?
-            GDALSetRasterColorInterpretation(hBand, <GDALColorInterp>1)
+            GDALSetRasterColorInterpretation(hBand, GCI_PaletteIndex)
             GDALSetRasterColorTable(hBand, hTable)
 
         finally:

tests/test_colormap.py

@@ -1,27 +1,26 @@
+"""Colormap tests."""
+
 import rasterio
+from rasterio.enums import ColorInterp
 
 
-def test_write_colormap_warn(tmpdir, recwarn):
-    with rasterio.open('tests/data/shade.tif') as src:
-        profile = src.meta
-    tiffname = str(tmpdir.join('foo.tif'))
-    with rasterio.open(tiffname, 'w', **profile) as dst:
-        dst.write_colormap(1, {0: (255, 0, 0, 255), 255: (0, 0, 0, 0)})
-
-
-def test_write_colormap(tmpdir):
-    with rasterio.open('tests/data/shade.tif') as src:
+def test_write_colormap(tmp_path):
+    with rasterio.open("tests/data/shade.tif") as src:
         shade = src.read(1)
-        meta = src.meta
-    tiffname = str(tmpdir.join('foo.png'))
-    meta['driver'] = 'PNG'
-    with rasterio.open(tiffname, 'w', **meta) as dst:
+        profile = src.profile
+
+    profile["driver"] = "PNG"
+
+    with rasterio.open(tmp_path / "test.tif", "w", **profile) as dst:
         dst.write(shade, indexes=1)
         dst.write_colormap(1, {0: (255, 0, 0, 255), 255: (0, 0, 0, 0)})
+        assert dst.colorinterp == (ColorInterp.palette,)
         cmap = dst.colormap(1)
         assert cmap[0] == (255, 0, 0, 255)
         assert cmap[255] == (0, 0, 0, 0)
-    with rasterio.open(tiffname) as src:
+
+    with rasterio.open(tmp_path / "test.tif") as src:
+        assert src.colorinterp == (ColorInterp.palette,)
         cmap = src.colormap(1)
         assert cmap[0] == (255, 0, 0, 255)
         assert cmap[255] == (0, 0, 0, 0)