DAS-2276: Refactor prepare_raster_for_writing.

Moves the palettize_raster function inside again.

hybig/browse.py

@@ -181,11 +181,9 @@ def create_browse_imagery(
                     f'incorrect number of bands for image: {rio_in_array.rio.count}'
                 )
 
-            raster, color_map = prepare_raster_for_writing(raster, output_driver)
-
-            if rio_in_array.rio.count == 1 and output_driver == 'PNG':
-                # we only paletize single band input data
-                raster, color_map = palettize_raster(raster)
+            raster, color_map = prepare_raster_for_writing(
+                raster, output_driver, rio_in_array.rio.count
+            )
 
             grid_parameters = get_target_grid_parameters(message, rio_in_array)
             grid_parameter_list, tile_locators = create_tiled_output_parameters(
@@ -360,12 +358,21 @@ def image_driver(mime: str) -> str:
 
 
 def prepare_raster_for_writing(
-    raster: ndarray, driver: str
+    raster: ndarray,
+    driver: str,
+    input_bands: int,
 ) -> tuple[ndarray, dict | None]:
     """Remove alpha layer if writing a jpeg."""
+    color_map = None
     if driver == 'JPEG' and raster.shape[0] == 4:
         raster = raster[0:3, :, :]
-    return raster, None
+        return raster, color_map
+
+    if input_bands == 1 and driver == 'PNG':
+        # we only paletize single band input data
+        raster, color_map = palettize_raster(raster)
+
+    return raster, color_map
 
 
 def palettize_raster(raster: ndarray) -> tuple[ndarray, dict]:

tests/test_service/test_adapter.py

@@ -460,7 +460,7 @@ def move_tif(*args, **kwargs):
             'count': 3,
         }
         raster = convert_mulitband_to_raster(rio_data_array)
-        raster, color_map = prepare_raster_for_writing(raster, 'PNG')
+        raster, color_map = prepare_raster_for_writing(raster, 'PNG', 3)
 
         dest = np.full(
             (expected_params['height'], expected_params['width']),

tests/unit/test_browse.py

@@ -655,38 +655,57 @@ def test_convert_5_multiband_to_raster(self):
 
     def test_prepare_raster_for_writing_jpeg_3band(self):
         raster = self.random.integers(255, size=(3, 5, 6))
+        count = 'irrelevant'
         driver = 'JPEG'
         expected_raster = np.copy(raster)
         expected_color_map = None
 
-        actual_raster, actual_color_map = prepare_raster_for_writing(raster, driver)
+        actual_raster, actual_color_map = prepare_raster_for_writing(
+            raster, driver, count
+        )
         self.assertEqual(expected_color_map, actual_color_map)
         np.testing.assert_array_equal(expected_raster, actual_raster, strict=True)
 
     def test_prepare_raster_for_writing_jpeg_4band(self):
         raster = self.random.integers(255, size=(4, 7, 8))
         driver = 'JPEG'
+        count = 'irrelevant'
         expected_raster = np.copy(raster[0:3, :, :])
         expected_color_map = None
-        actual_raster, actual_color_map = prepare_raster_for_writing(raster, driver)
+        actual_raster, actual_color_map = prepare_raster_for_writing(
+            raster, driver, count
+        )
         self.assertEqual(expected_color_map, actual_color_map)
         np.testing.assert_array_equal(expected_raster, actual_raster, strict=True)
 
     @patch('hybig.browse.palettize_raster')
     def test_prepare_raster_for_writing_png_4band(self, palettize_mock):
         raster = self.random.integers(255, size=(4, 7, 8))
         driver = 'PNG'
+        count = 'not 1'
 
-        expected, _ = prepare_raster_for_writing(raster, driver)
+        expected, _ = prepare_raster_for_writing(raster, driver, count)
         np.testing.assert_array_equal(raster, expected, strict=True)
+        palettize_mock.assert_not_called()
 
     @patch('hybig.browse.palettize_raster')
     def test_prepare_raster_for_writing_png_3band(self, palettize_mock):
         raster = self.random.integers(255, size=(3, 7, 8))
         driver = 'PNG'
+        count = 'not 1'
 
-        expected, _ = prepare_raster_for_writing(raster, driver)
+        expected, _ = prepare_raster_for_writing(raster, driver, count)
         np.testing.assert_array_equal(raster, expected, strict=True)
+        palettize_mock.assert_not_called()
+
+    @patch('hybig.browse.palettize_raster')
+    def test_prepare_1band_raster_for_writing_png(self, palettize_mock):
+        raster = self.random.integers(255, size=(1, 7, 8))
+        driver = 'PNG'
+        count = 1
+        palettize_mock.return_value = (None, None)
+        expected, _ = prepare_raster_for_writing(raster, driver, count)
+        palettize_mock.assert_called_with(raster)
 
     @patch('hybig.browse.Image')
     @patch('hybig.browse.get_color_map_from_image')