load TUNITS directly

spacetelescope · Dec 4, 2023 · 5b81bc4 · 5b81bc4
1 parent d3d37c4
commit 5b81bc4
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/src/stdatamodels/jwst/datamodels/multispec.py b/src/stdatamodels/jwst/datamodels/multispec.py
@@ -1,5 +1,6 @@
 from .model_base import JwstDataModel
 from .spec import SpecModel
+from astropy.io import fits
 
 
 __all__ = ['MultiSpecModel']
@@ -57,3 +58,13 @@ def __init__(self, init=None, **kwargs):
             return
 
         super(MultiSpecModel, self).__init__(init=init, **kwargs)
+
+        try:
+            if init[1].name == 'EXTRACT1D':
+                for key in init[1].header.keys():
+                    if 'TUNIT' in key:
+                        col = int(key.split('TUNIT')[1]) - 1
+                        for spec in self.spec:
+                            spec.spec_table.columns[col].unit = init[1].header[key]
+        except (AttributeError, IndexError) as e:
+            pass