Add colormap API to interface and tests

Author: mwcraig

src/astro_image_display_api/interface_definition.py

@@ -12,7 +12,23 @@
 
 # Allowed locations for cursor display
 ALLOWED_CURSOR_LOCATIONS = ('top', 'bottom', None)
-
+MINIMUM_REQUIRED_COLORMAPS = (
+    'gray',
+    'viridis',
+    'plasma',
+    'inferno',
+    'magma',
+    'purple-blue',
+    'yellow-green-blue',
+    'yellow-orange-red',
+    'red-purple',
+    'blue-green',
+    'hot',
+    'red-blue',
+    'red-yellow-blue',
+    'purple-orange'
+    'purple-green',
+)
 
 __all__ = [
     'ImageViewerInterface',
@@ -31,6 +47,9 @@ class ImageViewerInterface(Protocol):
     # Allowed locations for cursor display
     ALLOWED_CURSOR_LOCATIONS: tuple = ALLOWED_CURSOR_LOCATIONS
 
+    # Required colormaps for the viewer
+    MINIMUM_REQUIRED_COLORMAPS: tuple[str, ...] = MINIMUM_REQUIRED_COLORMAPS
+
     # The methods, grouped loosely by purpose
 
     # Method for loading image data
@@ -130,6 +149,69 @@ def get_stretch(self, image_label: str | None = None) -> BaseStretch:
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def set_colormap(self, map_name: str, image_label: str | None = None) -> None:
+        """
+        Set the colormap for the image specified by image_label.
+
+        Parameters
+        ----------
+        map_name : str
+            The name of the colormap to set. This should be a valid
+            colormap name from Matplotlib; not all backends will support
+            all colormaps, so the viewer should handle errors gracefully.
+            The case of the `map_name` is not important.
+        image_label : str, optional
+            The label of the image to set the colormap for. If not given and there is
+            only one image loaded, the colormap for that image is set. If there are
+            multiple images and no label is provided, an error is raised.
+
+        Raises
+        ------
+        ValueError
+            If the `map_name` is not a valid colormap name or if the `image_label`
+            is not provided when there are multiple images loaded.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_colormap(self, image_label: str | None = None) -> str:
+        """
+        Get the current colormap for the image.
+
+        Parameters
+        ----------
+        image_label : str, optional
+            The label of the image to get the colormap for. If not given and there is
+            only one image loaded, the colormap for that image is returned. If there are
+            multiple images and no label is provided, an error is raised.
+
+        Returns
+        -------
+        map_name : str
+            The name of the current colormap.
+
+        Raises
+        ------
+        ValueError
+            If the `image_label` is not provided when there are multiple images loaded or if
+            the `image_label` does not correspond to a loaded image.
+        """
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def colormap_options(self) -> list[str]:
+        """
+        Get the list of available colormaps.
+
+        Returns
+        -------
+        list of str
+            A list of available colormap names.
+        """
+        raise NotImplementedError
+
     # Saving contents of the view and accessing the view
     @abstractmethod
     def save(self, filename: str | os.PathLike, overwrite: bool = False) -> None:

src/astro_image_display_api/widget_api_test.py

@@ -7,11 +7,13 @@
 from astropy.coordinates import SkyCoord
 from astropy.io import fits
 from astropy.nddata import CCDData, NDData
-from astropy.table import Table, vstack
+from astropy.table import Table
 from astropy import units as u
 from astropy.wcs import WCS
 from astropy.visualization import AsymmetricPercentileInterval, BaseInterval, BaseStretch, LogStretch, ManualInterval
 
+from .interface_definition import ImageViewerInterface
+
 __all__ = ['ImageWidgetAPITest']
 
 
@@ -716,15 +718,48 @@ def test_stretch_cuts_errors(self, data):
         with pytest.raises(ValueError, match='[Ii]mage label.*not found'):
             self.image.set_cuts((10, 100), image_label='not a valid label')
 
-    @pytest.mark.skip(reason="Not clear whether colormap is part of the API")
-    def test_colormap(self):
-        cmap_desired = 'gray'
+    def test_colormap_options(self):
         cmap_list = self.image.colormap_options
-        assert len(cmap_list) > 0 and cmap_desired in cmap_list
+        assert set(ImageViewerInterface.MINIMUM_REQUIRED_COLORMAPS) <= set(cmap_list)
+        assert set(self.image.MINIMUM_REQUIRED_COLORMAPS) == set(ImageViewerInterface.MINIMUM_REQUIRED_COLORMAPS)
+
+    def test_set_get_colormap(self, data):
+        # Check setting and getting with a single image label.
+        self.image.load_image(data, image_label='test')
+        cmap_desired = 'gray'
         self.image.set_colormap(cmap_desired)
+        assert self.image.get_colormap() == cmap_desired
+
+        # Check that the colormap can be set with an image label
+        new_cmap = "viridis"
+        self.image.set_colormap(new_cmap, image_label='test')
+        assert self.image.get_colormap(image_label='test') == new_cmap
+
+    def test_set_colormap_errors(self, data):
+        # Check that setting a colormap raises an error if the colormap
+        # is not in the list of allowed colormaps.
+        self.image.load_image(data, image_label='test')
+
+        with pytest.raises(ValueError, match='[Ii]nvalid colormap'):
+            self.image.set_colormap('not a valid colormap')
+
+        # Check that getting a colormap for an image label that does not exist
+        with pytest.raises(ValueError, match='[Ii]mage label.*not found'):
+            self.image.get_colormap(image_label='not a valid label')
+
+        # Check that setting a colormap without an image label fails
+        # when there is more than one image label
+        self.image.load_image(data, image_label='another test')
+        with pytest.raises(ValueError, match='Multiple image labels defined'):
+            self.image.set_colormap('gray')
+
+        # Same for getting the colormap without an image label
+        with pytest.raises(ValueError, match='Multiple image labels defined'):
+            self.image.get_colormap()
 
     def test_cursor(self):
         assert self.image.cursor in self.image.ALLOWED_CURSOR_LOCATIONS
+        assert set(ImageViewerInterface.ALLOWED_CURSOR_LOCATIONS) == set(self.image.ALLOWED_CURSOR_LOCATIONS)
         with pytest.raises(self.cursor_error_classes):
             self.image.cursor = 'not a valid option'
         self.image.cursor = 'bottom'