From 57f03fff055b83309bb737b5a0a1fe7d39105a81 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Fri, 3 Nov 2023 14:24:03 -0400
Subject: [PATCH] consolidate has_wcs logic

---
 jdaviz/app.py                  | 12 ++++--------
 jdaviz/configs/imviz/helper.py |  3 ++-
 jdaviz/core/helpers.py         | 10 +---------
 jdaviz/utils.py                |  9 +++++++++
 4 files changed, 16 insertions(+), 18 deletions(-)

diff --git a/jdaviz/app.py b/jdaviz/app.py
index 3145931de7..3b8dfef2ff 100644
--- a/jdaviz/app.py
+++ b/jdaviz/app.py
@@ -14,7 +14,6 @@
 from astropy.io import fits
 from astropy.time import Time
 from astropy.utils.decorators import deprecated
-from astropy.wcs.wcsapi import BaseHighLevelWCS
 from echo import CallbackProperty, DictCallbackProperty, ListCallbackProperty
 from ipygoldenlayout import GoldenLayout
 from ipysplitpanes import SplitPanes
@@ -58,7 +57,7 @@
 from jdaviz.core.registries import (tool_registry, tray_registry, viewer_registry,
                                     data_parser_registry)
 from jdaviz.core.tools import ICON_DIR
-from jdaviz.utils import SnackbarQueue, alpha_index, MultiMaskSubsetState
+from jdaviz.utils import SnackbarQueue, alpha_index, data_has_valid_wcs, MultiMaskSubsetState
 from ipypopout import PopoutButton
 
 __all__ = ['Application', 'ALL_JDAVIZ_CONFIGS']
@@ -554,11 +553,8 @@ def _change_reference_data(self, new_refdata_label, viewer_id=None):
             )
 
         # only re-center the viewer if all data layers have WCS:
-        data_has_wcs = [
-            hasattr(d, 'coords') and isinstance(d.coords, BaseHighLevelWCS)
-            for d in viewer.data()
-        ]
-        if all(data_has_wcs):
+        has_wcs_per_data = [data_has_valid_wcs(d) for d in viewer.data()]
+        if all(has_wcs_per_data):
             # re-center the viewer on previous location.
             viewer.center_on(sky_cen)
 
@@ -2391,7 +2387,7 @@ def _expose_meta(key):
             'locked': False,
             'ndims': data.ndim,
             'type': typ,
-            'has_wcs': wcsaxes is not None,
+            'has_wcs': data_has_valid_wcs(data),
             'meta': {k: v for k, v in data.meta.items() if _expose_meta(k)},
             'children': []}
 
diff --git a/jdaviz/configs/imviz/helper.py b/jdaviz/configs/imviz/helper.py
index b6efc077e7..b3b19ef099 100644
--- a/jdaviz/configs/imviz/helper.py
+++ b/jdaviz/configs/imviz/helper.py
@@ -17,6 +17,7 @@
 from jdaviz.configs.imviz.wcs_utils import (
     _get_rotated_nddata_from_label, get_compass_info
 )
+from jdaviz.utils import data_has_valid_wcs
 
 __all__ = ['Imviz', 'link_image_data']
 
@@ -198,7 +199,7 @@ def load_data(self, data, data_label=None, do_link=True, show_in_viewer=True, **
                 applied_labels.append(label)
                 applied_visible.append(True)
                 layer_is_wcs_only.append(data.meta.get(self.app._wcs_only_label, False))
-                layer_has_wcs.append(hasattr(data.coords, 'pixel_to_world'))
+                layer_has_wcs.append(data_has_valid_wcs(data))
 
         if show_in_viewer is True:
             show_in_viewer = f"{self.app.config}-0"
diff --git a/jdaviz/core/helpers.py b/jdaviz/core/helpers.py
index faafa0afd0..5119f80b31 100644
--- a/jdaviz/core/helpers.py
+++ b/jdaviz/core/helpers.py
@@ -13,7 +13,6 @@
 
 import numpy as np
 import astropy.units as u
-from astropy.wcs.wcsapi import BaseHighLevelWCS
 from astropy.nddata import CCDData, StdDevUncertainty
 from regions.core.core import Region
 from glue.core import HubListener
@@ -28,6 +27,7 @@
 from jdaviz.app import Application
 from jdaviz.core.events import SnackbarMessage, ExitBatchLoadMessage
 from jdaviz.core.template_mixin import show_widget
+from jdaviz.utils import data_has_valid_wcs
 
 __all__ = ['ConfigHelper', 'ImageConfigHelper']
 
@@ -900,14 +900,6 @@ def _delete_all_regions(self):
             self.app.data_collection.remove_subset_group(subset_grp)
 
 
-def data_has_valid_wcs(data, ndim=None):
-    """Check if given glue Data has WCS that is compatible with APE 14."""
-    status = hasattr(data, 'coords') and isinstance(data.coords, BaseHighLevelWCS)
-    if ndim is not None:
-        status = status and data.coords.world_n_dim == ndim
-    return status
-
-
 def _next_subset_num(label_prefix, subset_groups):
     """Assumes ``prefix i`` format.
     Does not go back and fill in lower but available numbers. This is consistent with Glue.
diff --git a/jdaviz/utils.py b/jdaviz/utils.py
index 9be5230092..5a9ab3bc14 100644
--- a/jdaviz/utils.py
+++ b/jdaviz/utils.py
@@ -6,6 +6,7 @@
 
 from astropy.io import fits
 from astropy.utils import minversion
+from astropy.wcs.wcsapi import BaseHighLevelWCS
 from ipyvue import watch
 from glue.core.exceptions import IncompatibleAttribute
 from glue.config import settings
@@ -216,6 +217,14 @@ def alpha_index(index):
         return chr(97 + index//26 - 1) + chr(97 + index % 26)
 
 
+def data_has_valid_wcs(data, ndim=None):
+    """Check if given glue Data has WCS that is compatible with APE 14."""
+    status = hasattr(data, 'coords') and isinstance(data.coords, BaseHighLevelWCS)
+    if ndim is not None:
+        status = status and data.coords.world_n_dim == ndim
+    return status
+
+
 def standardize_metadata(metadata):
     """Standardize given metadata so it can be viewed in
     Metadata Viewer plugin. The input can be plain