Handle checking GWCS for spectral subset, fix codestyle

jdaviz/configs/cubeviz/plugins/tests/test_parsers.py

@@ -212,8 +212,8 @@ def test_numpy_cube(cubeviz_helper):
 def test_loading_with_mask(cubeviz_helper):
     # This also tests that spaxel is converted to pix**2
     custom_spec = Spectrum(flux=[[[20, 1], [9, 1]], [[3, 1], [6, np.nan]]] * u.Unit("erg / Angstrom s cm**2 spaxel"),  # noqa
-                             spectral_axis=[1, 2]*u.AA,
-                             mask=[[[1, 0], [0, 0]], [[0, 0], [0, 0]]])
+                           spectral_axis=[1, 2]*u.AA,
+                           mask=[[[1, 0], [0, 0]], [[0, 0], [0, 0]]])
     cubeviz_helper.load_data(custom_spec)
 
     uc = cubeviz_helper.plugins['Unit Conversion']

jdaviz/utils.py

@@ -14,6 +14,8 @@
 from astropy.units import Quantity
 from astropy import units as u
 from astroquery.mast import Observations, conf
+from gwcs import WCS as gwcs
+from gwcs.coordinate_frames import CompositeFrame, SpectralFrame
 from matplotlib import colors as mpl_colors
 import matplotlib.cm as cm
 from photutils.utils import make_random_cmap
@@ -624,6 +626,17 @@ def get_subset_type(subset):
                 # subset type:
                 continue
 
+            # check for spectral coordinate in GWCS by looking for SpectralFrame
+            if isinstance(ss_data.coords, gwcs):
+                if isinstance(ss_data.coords, SpectralFrame):
+                    return 'spectral'
+                elif isinstance(ss_data.coords, CompositeFrame):
+                    if np.any([isinstance(frame, SpectralFrame) for frame in
+                               ss_data.coords.output_frame.frames]):
+                        return 'spectral'
+                else:
+                    continue
+
             # check for a spectral coordinate in FITS WCS:
             wcs_coords = (
                 ss_data.coords.wcs.ctype if hasattr(ss_data.coords, 'wcs')