Skip to content
Back to top

# Source code for gwcs.wcs._step

    import warnings
    from typing import NamedTuple, TypeAlias, Union

    from astropy.modeling.core import Model

    from gwcs.coordinate_frames import CoordinateFrame, EmptyFrame

    __all__ = [
        "IndexedStep",
        "Step",
        "StepTuple",
    ]

    StepTuple: TypeAlias = tuple[CoordinateFrame, Union[Model, None]]  # noqa: UP007

    class Step:
        """
        Represents a ``step`` in the WCS pipeline.

        Parameters
        ----------
        frame : `~gwcs.coordinate_frames.CoordinateFrame`
            A gwcs coordinate frame object.
        transform : `~astropy.modeling.Model` or None
            A transform from this step's frame to next step's frame.
            The transform of the last step should be `None`.
        """

        def __init__(self, frame: str | CoordinateFrame | None, transform=None):
            # Allow for a string to be passed in for the frame but be turned into a
            # frame object
            self.frame = (
                frame if isinstance(frame, CoordinateFrame) else EmptyFrame(name=frame)
            )
            self.transform = transform

        @property
        def frame(self):
            return self._frame

        @frame.setter
        def frame(self, val):
            if not isinstance(val, CoordinateFrame | str):
                msg = '"frame" should be an instance of CoordinateFrame or a string.'
                raise TypeError(msg)

            self._frame = val

        @property
        def transform(self):
            return self._transform

        @transform.setter
        def transform(self, val):
            if val is not None and not isinstance(val, (Model)):
                msg = '"transform" should be an instance of astropy.modeling.Model.'
                raise TypeError(msg)
            self._transform = val

        @property
        def frame_name(self):
            if isinstance(self.frame, str):
                return self.frame
            return self.frame.name

        def __getitem__(self, ind):
            warnings.warn(
                "Indexing a WCS.pipeline step is deprecated. "
                "Use the `frame` and `transform` attributes instead.",
                DeprecationWarning,
                stacklevel=2,
            )
            if ind not in (0, 1):
                msg = "Allowed inices are 0 (frame) and 1 (transform)."
                raise IndexError(msg)
            if ind == 0:
                return self.frame
            return self.transform

        def __str__(self):
            return (
                f"{self.frame_name}\t "
                f"{getattr(self.transform, 'name', 'None') or type(self.transform).__name__}"  # noqa: E501
            )

        def __repr__(self):
            return (
                f"Step(frame={self.frame_name}, "
                f"transform={getattr(self.transform, 'name', 'None') or type(self.transform).__name__})"  # noqa: E501
            )

        def copy(self):
            return Step(self.frame, self.transform)

    class IndexedStep(NamedTuple):
        """
        Class to handle a step and its index in the pipeline.
        """

        idx: int
        step: Step