More implementation WIP dummy

Author: mwcraig

src/astro_image_display_api/dummy_viewer.py

@@ -1,4 +1,6 @@
 import os
+from collections import defaultdict
+from copy import deepcopy
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any
@@ -12,6 +14,13 @@
 
 from .interface_definition import ImageViewerInterface
 
+@dataclass
+class CatalogInfo:
+    """
+    A named tuple to hold information about a catalog.
+    """
+    style: dict[str, Any] = field(default_factory=dict)
+    data: Table | None = None
 
 @dataclass
 class ImageViewer:
@@ -28,9 +37,8 @@ class ImageViewer:
     stretch_options: tuple = ("linear", "log", "sqrt")
     autocut_options: tuple = ("minmax", "zscale", "asinh", "percentile", "histogram")
     _cursor: str = ImageViewerInterface.ALLOWED_CURSOR_LOCATIONS[0]
-    _markers: dict[str, dict] = field(default_factory=dict)
-    _catalogs: dict[str, Table] = field(default_factory=dict)
-    _catalog_names: list[str] = field(default_factory=list)
+    # marker style will be stored in catalog table metadata
+    _catalogs: dict[str, CatalogInfo] = field(default_factory=dict)
     _default_marker_style: dict[str, Any] = field(default_factory=dict)
     _cuts: str | tuple[float, float] = (0, 1)
     _stretch: str = "linear"
@@ -47,24 +55,32 @@ def __post_init__(self):
         # This is a dictionary of marker sets. The keys are the names of the
         # marker sets, and the values are the tables containing the markers.
         self._default_marker_style = dict(shape="circle", color="yellow", size=10)
-        self._markers[None] = self._default_marker_style.copy()
+        self._catalogs = defaultdict(CatalogInfo)
+        self._catalogs[None].data = None
+        self._catalogs[None].style = self._default_marker_style.copy()
 
     def _user_catalog_labels(self) -> list[str]:
         """
         Get the user-defined catalog labels.
         """
-        return [label for label in self._markers if label is not None]
+        return [label for label in self._catalogs if label is not None]
 
-    def _resolve_catalog_label(self, catalog_label: str | None) -> list[str]:
+    def _resolve_catalog_label(self, catalog_label: str | None) -> str:
+        """
+        Figure out the catalog label if the user did not specify one. This
+        is needed so that the user gets what they expect in the simple case
+        where there is only one catalog loaded. In that case the user may
+        or may not have actually specified a catalog label.
+        """
         user_keys = self._user_catalog_labels()
         if catalog_label is None:
             match len(user_keys):
                 case 0:
-                    # No user-defined styles, so return the default style
+                    # No user-defined catalog labels, so return the default label.
                     catalog_label = None
                 case 1:
-                    # The user must have set a style, so return that instead of
-                    # the default style, which live in the key None.
+                    # The user must have loaded a catalog, so return that instead of
+                    # the default label, which live in the key None.
                     catalog_label = user_keys[0]
                 case _:
                     raise ValueError(
@@ -128,7 +144,7 @@ def get_catalog_style(self, catalog_label=None) -> dict[str, dict[str, Any]]:
         """
         catalog_label = self._resolve_catalog_label(catalog_label)
 
-        style = self._markers[catalog_label]
+        style = self._catalogs[catalog_label].style
         style["catalog_label"] = catalog_label
         return style
 
@@ -160,7 +176,12 @@ def set_catalog_style(
         color = color if color else self._default_marker_style["color"]
         size = size if size else self._default_marker_style["size"]
 
-        self._markers[catalog_label] = {
+        catalog_label = self._resolve_catalog_label(catalog_label)
+
+        if self._catalogs[catalog_label].data is None:
+            raise ValueError("Must load a catalog before setting a catalog style.")
+
+        self._catalogs[catalog_label].style = {
             "shape": shape,
             "color": color,
             "size": size,
@@ -242,8 +263,9 @@ def save(self, filename: str | os.PathLike, overwrite: bool = False) -> None:
 
     # Marker-related methods
     def load_catalog(self, table: Table, x_colname: str = 'x', y_colname: str = 'y',
-                    skycoord_colname: str = 'coord', use_skycoord: bool = False,
-                    catalog_label: str | None = None) -> None:
+                    skycoord_colname: str = 'coord', use_skycoord: bool = True,
+                    catalog_label: str | None = None,
+                    catalog_style: dict | None = None) -> None:
         """
         Add markers to the image.
 
@@ -263,7 +285,7 @@ def load_catalog(self, table: Table, x_colname: str = 'x', y_colname: str = 'y',
             is ``'coord'``.
         use_skycoord : bool, optional
             If `True`, the ``skycoord_colname`` column will be used to
-            get the marker positions. Default is `False`.
+            get the marker positions.
         catalog_label : str, optional
             The name of the marker set to use. If not given, a unique
             name will be generated.
@@ -273,32 +295,39 @@ def load_catalog(self, table: Table, x_colname: str = 'x', y_colname: str = 'y',
         except KeyError:
             coords = None
 
-        if use_skycoord:
-            if self._wcs is not None:
+        try:
+            xy = (table[x_colname], table[y_colname])
+        except KeyError:
+            xy = None
+
+        to_add = deepcopy(table)
+        if xy is None:
+            if self._wcs is not None and coords is not None:
                 x, y = self._wcs.world_to_pixel(coords)
+                to_add[x_colname] = x
+                to_add[y_colname] = y
             else:
-                raise ValueError("WCS is not set. Cannot convert to pixel coordinates.")
-        else:
-            x = table[x_colname]
-            y = table[y_colname]
-
-            if not coords and self._wcs is not None:
-                coords = self._wcs.pixel_to_world(x, y)
+                to_add[x_colname] = to_add[y_colname] = None
 
-        to_add = Table(
-            dict(
-                x=x,
-                y=y,
-                coord=coords if coords else [None] * len(x),
-            )
-        )
+        if coords is None:
+            if use_skycoord and self._wcs is None:
+                raise ValueError("WCS is not set. Cannot convert to pixel coordinates.")
+            elif xy is not None and self._wcs is not None:
+                # If we have xy coordinates, convert them to sky coordinates
+                coords = self._wcs.pixel_to_world(xy[0], xy[1])
+                to_add[skycoord_colname] = coords
+            else:
+                to_add[skycoord_colname] = None
 
         catalog_label = self._resolve_catalog_label(catalog_label)
-        if catalog_label in self._catalogs:
-            marker_table = self._markers[catalog_label]
-            self._markers[catalog_label] = vstack([marker_table, to_add])
+        if (
+            catalog_label in self._catalogs
+            and self._catalogs[catalog_label].data is not None
+        ):
+            old_table = self._catalogs[catalog_label].data
+            self._catalogs[catalog_label].data = vstack([old_table, to_add])
         else:
-            self._markers[catalog_label] = to_add
+            self._catalogs[catalog_label].data = to_add
 
     def remove_catalog(self, catalog_label: str | None = None) -> None:
         """
@@ -310,19 +339,24 @@ def remove_catalog(self, catalog_label: str | None = None) -> None:
             The name of the marker set to remove. If the value is ``"*"``,
             then all markers will be removed.
         """
-        if isinstance(catalog_label, str):
-            if catalog_label in self._markers:
-                del self._markers[catalog_label]
-            elif catalog_label == "*":
-                self._markers = {}
-            else:
-                raise ValueError(f"Marker name {catalog_label} not found.")
-        elif isinstance(catalog_label, list):
-            for name in catalog_label:
-                if name in self._markers:
-                    del self._markers[name]
-                else:
-                    raise ValueError(f"Marker name {name} not found.")
+        if isinstance(catalog_label, list):
+            raise ValueError(
+                "Cannot remove multiple catalogs from a list. Please specify "
+                "a single catalog label or use '*' to remove all catalogs."
+            )
+        elif catalog_label == "*":
+            # If the user wants to remove all catalogs, we reset the
+            # catalogs dictionary to an empty one.
+            self._catalogs = defaultdict(CatalogInfo)
+            return
+
+        # Special cases are done, so we can resolve the catalog label
+        catalog_label = self._resolve_catalog_label(catalog_label)
+
+        try:
+            del self._catalogs[catalog_label]
+        except KeyError:
+            raise ValueError(f"Marker name {catalog_label} not found.")
 
     def get_catalog(self, x_colname: str = 'x', y_colname: str = 'y',
                     skycoord_colname: str = 'coord',
@@ -331,7 +365,7 @@ def get_catalog(self, x_colname: str = 'x', y_colname: str = 'y',
         # duplicated here.
         catalog_label = self._resolve_catalog_label(catalog_label)
 
-        result = self._markers[catalog_label] if catalog_label in self._markers else Table(names=["x", "y", "coord", "marker name"])
+        result = self._catalogs[catalog_label].data if catalog_label in self._catalogs else Table(names=["x", "y", "coord", "marker name"])
 
         result.rename_columns(["x", "y", "coord"], [x_colname, y_colname, skycoord_colname])