From 996a5c8d683af8a5b327185bfe76fcc8b985662d Mon Sep 17 00:00:00 2001
From: "P. L. Lim" <2090236+pllim@users.noreply.github.com>
Date: Wed, 15 Feb 2023 13:48:37 -0500
Subject: [PATCH] Fix bug in cunit lookup

---
 jdaviz/configs/cubeviz/plugins/parsers.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/jdaviz/configs/cubeviz/plugins/parsers.py b/jdaviz/configs/cubeviz/plugins/parsers.py
index 8cb455367a..331d062459 100644
--- a/jdaviz/configs/cubeviz/plugins/parsers.py
+++ b/jdaviz/configs/cubeviz/plugins/parsers.py
@@ -149,8 +149,10 @@ def _return_spectrum_with_correct_units(flux, wcs, metadata, data_type, target_w
                 continue
             hdr = hdulist[ext].header
             # The WCS could be swapped or unswapped.
-            for cunit_key in ('CUNIT3', 'CUNIT1'):
-                if not found_target and cunit_key in hdr and 'WAVE' in hdr[cunit_key]:
+            for cunit_num in (3, 1):
+                cunit_key = f"CUNIT{cunit_num}"
+                ctype_key = f"CTYPE{cunit_num}"
+                if not found_target and cunit_key in hdr and 'WAVE' in hdr[ctype_key]:
                     target_wave_unit = u.Unit(hdr[cunit_key])
                     found_target = True
                     break