JP-3860: Apply code style rules to assign_wcs

.pre-commit-config.yaml

@@ -37,7 +37,6 @@ repos:
       - id: numpydoc-validation
         exclude: |
           (?x)^(
-            jwst/assign_wcs/.* |
             jwst/associations/.* |
             jwst/background/.* |
             jwst/coron/.* |

.ruff.toml

@@ -19,7 +19,6 @@ quote-style = "double"
 indent-style = "space"
 docstring-code-format = true
 exclude = [
-    "jwst/assign_wcs/**.py",
     "jwst/associations/**.py",
     "jwst/background/**.py",
     "jwst/coron/**.py",
@@ -118,7 +117,6 @@ ignore-fully-untyped = true  # Turn off annotation checking for fully untyped co
 "jwst/associations/tests*" = [
     "F841", # unused variable
 ]
-"jwst/assign_wcs/**.py" = ["D", "N", "A", "ARG", "B", "C4", "ICN", "INP", "ISC", "LOG", "NPY", "PGH", "PTH", "S", "SLF", "SLOT", "T20", "TRY", "UP", "YTT", "E501"]
 "jwst/associations/**.py" = ["D", "N", "A", "ARG", "B", "C4", "ICN", "INP", "ISC", "LOG", "NPY", "PGH", "PTH", "S", "SLF", "SLOT", "T20", "TRY", "UP", "YTT", "E501"]
 "jwst/background/**.py" = ["D", "N", "A", "ARG", "B", "C4", "ICN", "INP", "ISC", "LOG", "NPY", "PGH", "PTH", "S", "SLF", "SLOT", "T20", "TRY", "UP", "YTT", "E501"]
 "jwst/coron/**.py" = ["D", "N", "A", "ARG", "B", "C4", "ICN", "INP", "ISC", "LOG", "NPY", "PGH", "PTH", "S", "SLF", "SLOT", "T20", "TRY", "UP", "YTT", "E501"]

jwst/assign_wcs/__init__.py

@@ -1,7 +1,15 @@
+"""Assign WCS information to JWST data models."""
+
 from .assign_wcs_step import AssignWcsStep
-from .nirspec import (nrs_wcs_set_input, nrs_ifu_wcs, get_spectral_order_wrange)
+from .nirspec import nrs_wcs_set_input, nrs_ifu_wcs, get_spectral_order_wrange
 from .niriss import niriss_soss_set_input
 from .util import update_fits_wcsinfo
 
-__all__ = ['AssignWcsStep', "nrs_wcs_set_input", "nrs_ifu_wcs", "get_spectral_order_wrange",
-           "niriss_soss_set_input", "update_fits_wcsinfo"]
+__all__ = [
+    "AssignWcsStep",
+    "nrs_wcs_set_input",
+    "nrs_ifu_wcs",
+    "get_spectral_order_wrange",
+    "niriss_soss_set_input",
+    "update_fits_wcsinfo",
+]

jwst/assign_wcs/assign_wcs.py

@@ -1,8 +1,13 @@
 import logging
 import importlib
 from gwcs.wcs import WCS
-from .util import (update_s_region_spectral, update_s_region_imaging,
-                   update_s_region_nrs_ifu, update_s_region_mrs)
+from .util import (
+    update_s_region_spectral,
+    update_s_region_imaging,
+    update_s_region_nrs_ifu,
+    update_s_region_mrs,
+    update_s_region_lrs,
+)
 from ..lib.exposure_types import IMAGING_TYPES, SPEC_TYPES, NRS_LAMP_MODE_SPEC_TYPES
 from ..lib.dispaxis import get_dispersion_direction
 from ..lib.wcs_utils import get_wavelengths
@@ -14,91 +19,113 @@
 __all__ = ["load_wcs"]
 
 
-def load_wcs(input_model, reference_files={}, nrs_slit_y_range=None):
+def load_wcs(input_model, reference_files=None, nrs_slit_y_range=None):
     """
     Create a gWCS object and store it in ``Model.meta``.
 
     Parameters
     ----------
     input_model : `~jwst.datamodels.JwstDataModel`
-        The exposure.
+        The input data model.
     reference_files : dict
-        A dict {reftype: reference_file_name} containing all
-        reference files that apply to this exposure.
+        Mapping between reftype (keys) and reference file name (vals).
     nrs_slit_y_range : list
         The slit y-range for a Nirspec slit. The center is (0, 0).
+
+    Returns
+    -------
+    output_model : `~jwst.datamodels.JwstDataModel`
+        The data model with the WCS information in the meta attribute.
     """
-    if reference_files:
+    if reference_files is not None:
         for ref_type, ref_file in reference_files.items():
             if ref_file not in ["N/A", ""]:
                 reference_files[ref_type] = ref_file
             else:
                 reference_files[ref_type] = None
-    if not any(reference_files.values()):
+    if (reference_files is None) or (not any(reference_files.values())):
         log.critical("assign_wcs needs reference files to compute the WCS, none were passed")
         raise ValueError("assign_wcs needs reference files to compute the WCS, none were passed")
     instrument = input_model.meta.instrument.name.lower()
-    mod = importlib.import_module('.' + instrument, 'jwst.assign_wcs')
+    mod = importlib.import_module("." + instrument, "jwst.assign_wcs")
 
-    if input_model.meta.exposure.type.lower() in SPEC_TYPES or \
-       input_model.meta.instrument.lamp_mode.lower() in NRS_LAMP_MODE_SPEC_TYPES:
+    if (
+        input_model.meta.exposure.type.lower() in SPEC_TYPES
+        or input_model.meta.instrument.lamp_mode.lower() in NRS_LAMP_MODE_SPEC_TYPES
+    ):
         input_model.meta.wcsinfo.specsys = "BARYCENT"
-        input_model.meta.wcsinfo.dispersion_direction = \
-            get_dispersion_direction(
-                input_model.meta.exposure.type,
-                input_model.meta.instrument.grating,
-                input_model.meta.instrument.filter,
-                input_model.meta.instrument.pupil)
+        input_model.meta.wcsinfo.dispersion_direction = get_dispersion_direction(
+            input_model.meta.exposure.type,
+            input_model.meta.instrument.grating,
+            input_model.meta.instrument.filter,
+            input_model.meta.instrument.pupil,
+        )
 
-    if instrument.lower() == 'nirspec':
+    if instrument.lower() == "nirspec":
         pipeline = mod.create_pipeline(input_model, reference_files, slit_y_range=nrs_slit_y_range)
     else:
         pipeline = mod.create_pipeline(input_model, reference_files)
     # Initialize the output model as a copy of the input
     # Make the copy after the WCS pipeline is created in order to pass updates to the model.
     if pipeline is None:
-        input_model.meta.cal_step.assign_wcs = 'SKIPPED'
+        input_model.meta.cal_step.assign_wcs = "SKIPPED"
         log.warning("assign_wcs: SKIPPED")
         return input_model
-    else:
-        output_model = input_model.copy()
-        wcs = WCS(pipeline)
-        output_model.meta.wcs = wcs
-        output_model.meta.cal_step.assign_wcs = 'COMPLETE'
-        exclude_types = ['nrc_wfss', 'nrc_tsgrism', 'nis_wfss',
-                         'nrs_fixedslit', 'nrs_msaspec',
-                         'nrs_autowave', 'nrs_autoflat', 'nrs_lamp',
-                         'nrs_brightobj', 'nis_soss']
 
-        if output_model.meta.exposure.type.lower() not in exclude_types:
-            imaging_types = IMAGING_TYPES.copy()
-            imaging_types.update(['mir_lrs-fixedslit', 'mir_lrs-slitless'])
-            if output_model.meta.exposure.type.lower() in imaging_types:
-                try:
-                    update_s_region_imaging(output_model)
-                except Exception as exc:
-                    log.error("Unable to update S_REGION for type {}: {}".format(
-                        output_model.meta.exposure.type, exc))
-                else:
-                    log.info("assign_wcs updated S_REGION to {0}".format(
-                        output_model.meta.wcsinfo.s_region))
-                if output_model.meta.exposure.type.lower() == 'mir_lrs-slitless':
-                    output_model.wavelength = get_wavelengths(output_model)
-            elif output_model.meta.exposure.type.lower() == "nrs_ifu":
-                update_s_region_nrs_ifu(output_model, mod)
-            elif output_model.meta.exposure.type.lower() == 'mir_mrs':
-                update_s_region_mrs(output_model)
+    output_model = input_model.copy()
+    wcs = WCS(pipeline)
+    output_model.meta.wcs = wcs
+    output_model.meta.cal_step.assign_wcs = "COMPLETE"
+    exclude_types = [
+        "nrc_wfss",
+        "nrc_tsgrism",
+        "nis_wfss",
+        "nrs_fixedslit",
+        "nrs_msaspec",
+        "nrs_autowave",
+        "nrs_autoflat",
+        "nrs_lamp",
+        "nrs_brightobj",
+        "nis_soss",
+    ]
+
+    if output_model.meta.exposure.type.lower() not in exclude_types:
+        imaging_types = IMAGING_TYPES.copy()
+        imaging_types.update(["mir_lrs-slitless"])
+        imaging_lrs_types = ["mir_lrs-fixedslit"]
+        if output_model.meta.exposure.type.lower() in imaging_lrs_types:
+            # uses slits corners in V2, V3 that are read in from the
+            # lrs specwcs reference file
+            update_s_region_lrs(output_model, reference_files)
+        elif output_model.meta.exposure.type.lower() in imaging_types:
+            try:
+                update_s_region_imaging(output_model)
+            except Exception as exc:
+                log.error(
+                    f"Unable to update S_REGION for type {output_model.meta.exposure.type}: {exc}"
+                )
             else:
-                try:
-                    update_s_region_spectral(output_model)
-                except Exception as exc:
-                    log.info("Unable to update S_REGION for type {}: {}".format(
-                        output_model.meta.exposure.type, exc))
+                log.info(f"assign_wcs updated S_REGION to {output_model.meta.wcsinfo.s_region}")
+            if output_model.meta.exposure.type.lower() == "mir_lrs-slitless":
+                output_model.wavelength = get_wavelengths(output_model)
+        elif output_model.meta.exposure.type.lower() == "nrs_ifu":
+            update_s_region_nrs_ifu(output_model, mod)
+        elif output_model.meta.exposure.type.lower() == "mir_mrs":
+            update_s_region_mrs(output_model)
+        else:
+            try:
+                update_s_region_spectral(output_model)
+            except Exception as exc:
+                log.info(
+                    f"Unable to update S_REGION for type {output_model.meta.exposure.type}: {exc}"
+                )
 
-        # Store position of dithered pointing location in metadata for later spectral extraction
-        if output_model.meta.exposure.type.lower() == 'mir_lrs-fixedslit':
-            store_dithered_position(output_model)
-            log.debug(f"Storing dithered pointing location information:"
-                      f"{output_model.meta.dither.dithered_ra} {output_model.meta.dither.dithered_dec}")
+    # Store position of dithered pointing location in metadata for later spectral extraction
+    if output_model.meta.exposure.type.lower() == "mir_lrs-fixedslit":
+        store_dithered_position(output_model)
+        log.debug(
+            "Storing dithered pointing location information: "
+            f"{output_model.meta.dither.dithered_ra} {output_model.meta.dither.dithered_dec}"
+        )
     log.info("COMPLETED assign_wcs")
     return output_model