Merge pull request astropy#49 from mwcraig/add-colormap-appi

Add colormap api

Author: mwcraig

src/astro_image_display_api/dummy_viewer.py

@@ -10,7 +10,7 @@
 from astropy.coordinates import SkyCoord
 from astropy.nddata import CCDData, NDData
 from astropy.table import Table, vstack
-from astropy.units import Quantity, get_physical_type
+from astropy.units import Quantity
 from astropy.wcs import WCS
 from astropy.wcs.utils import proj_plane_pixel_scales
 from astropy.visualization import AsymmetricPercentileInterval, BaseInterval, BaseStretch, LinearStretch, ManualInterval
@@ -37,6 +37,7 @@ class ViewportInfo:
     largest_dimension: int | None = None
     stretch: BaseStretch | None = None
     cuts: BaseInterval | tuple[numbers.Real, numbers.Real] | None = None
+    colormap: str | None = None
 
 @dataclass
 class ImageViewer:
@@ -150,6 +151,22 @@ def set_cuts(self, value: tuple[numbers.Real, numbers.Real] | BaseInterval, imag
             raise ValueError(f"Image label '{image_label}' not found. Please load an image first.")
         self._images[image_label].cuts = self._cuts
 
+    def set_colormap(self, map_name: str, image_label: str | None = None) -> None:
+        image_label = self._resolve_image_label(image_label)
+        if image_label not in self._images:
+            raise ValueError(f"Image label '{image_label}' not found. Please load an image first.")
+        self._images[image_label].colormap = map_name
+
+    set_colormap.__doc__ = ImageViewerInterface.set_colormap.__doc__
+
+    def get_colormap(self, image_label: str | None = None) -> str:
+        image_label = self._resolve_image_label(image_label)
+        if image_label not in self._images:
+            raise ValueError(f"Image label '{image_label}' not found. Please load an image first.")
+        return self._images[image_label].colormap
+
+    get_colormap.__doc__ = ImageViewerInterface.get_colormap.__doc__
+
     @property
     def cursor(self) -> str:
         return self._cursor

src/astro_image_display_api/interface_definition.py

@@ -3,17 +3,14 @@
 import os
 
 from astropy.coordinates import SkyCoord
-from astropy.nddata import NDData
 from astropy.table import Table
 from astropy.units import Quantity
 from astropy.visualization import BaseInterval, BaseStretch
 
-from numpy.typing import ArrayLike
 
 # Allowed locations for cursor display
 ALLOWED_CURSOR_LOCATIONS = ('top', 'bottom', None)
 
-
 __all__ = [
     'ImageViewerInterface',
 ]
@@ -130,6 +127,58 @@ 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.
+        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.
+
+        .. _Matplotlib: https://matplotlib.org/stable/gallery/color/colormap_reference.html
+        """
+        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
+
     # 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,40 @@ 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):
+    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'
-        cmap_list = self.image.colormap_options
-        assert len(cmap_list) > 0 and cmap_desired in cmap_list
         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')
+
+        # 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'