noaa-goes-visualization/suvi/paths.py

"""Archive layout: filename grammar, slot keys, and root discovery.

The SUVI archive stores one FITS file per (satellite, wavelength, 4-minute slot):

    Data/goes16/l2/data/suvi-l2-ci171/2024/03/15/
        dr_suvi-l2-ci171_g16_s20240315T000000Z_e20240315T000400Z_v1-0-2.fits

Historically ``filter_FITS.py`` appended ``_f`` (passed) or ``_e`` (rejected) to the
stem to record its verdict.  Those labels now live in the SQLite index instead, but
the parser still recognises them so the migration can ingest and strip them.
"""

import datetime
import os
import re
from dataclasses import dataclass

# The six SUVI composite-image passbands, in Angstroms.
WAVELENGTHS = (94, 131, 171, 195, 284, 304)
# GOES-17 is deliberately absent: its SUVI data is excluded by puller_fits.py.
SATELLITES = (16, 18, 19)
# Nominal spacing between consecutive frames, in seconds.
CADENCE = 240

# Legacy verdict suffixes written into filenames by the pre-SQLite filter.
LEGACY_LABELS = ("f", "e")

_TIME_FORMAT = "%Y%m%dT%H%M%S"

# Anchored so that only an exact, well-formed basename matches.  The trailing label
# group repeats to absorb the double-suffix files ("..._f_f.fits") that earlier
# filter runs produced by re-processing an already-labelled file.
_FRAME_RE = re.compile(
    r"^dr_suvi-l2-ci(?P<wavelength>\d{3})"
    r"_g(?P<satellite>\d{2})"
    r"_s(?P<t_start>\d{8}T\d{6})Z"
    r"_e(?P<t_end>\d{8}T\d{6})Z"
    r"_v(?P<version>[0-9]+(?:-[0-9]+)*)"
    r"(?P<labels>(?:_[fe])*)"
    r"\.fits$"
)


def _parse_time(text):
    """Parse a filename timestamp token to a UTC Unix timestamp."""
    stamp = datetime.datetime.strptime(text, _TIME_FORMAT)
    return int(stamp.replace(tzinfo=datetime.timezone.utc).timestamp())


def _format_time(timestamp):
    stamp = datetime.datetime.fromtimestamp(timestamp, datetime.timezone.utc)
    return stamp.strftime(_TIME_FORMAT)


@dataclass(frozen=True)
class FrameName:
    """A parsed SUVI L2 composite-image filename."""

    satellite: int
    wavelength: int
    t_start: int
    t_end: int
    version: str
    #: Legacy verdict suffix, or None once the archive has been un-renamed.  A file
    #: labelled more than once keeps only the outermost (most recent) verdict.
    label: str | None = None

    @property
    def slot(self):
        """The key identifying what this frame is an observation of."""
        return (self.satellite, self.wavelength, self.t_start)

    def filename(self, label=None):
        """Render the filename, optionally with a legacy verdict suffix."""
        if label is not None and label not in LEGACY_LABELS:
            raise ValueError(f"Unknown label: {label!r}")
        suffix = f"_{label}" if label else ""
        return (
            f"dr_suvi-l2-ci{self.wavelength:03d}"
            f"_g{self.satellite:d}"
            f"_s{_format_time(self.t_start)}Z"
            f"_e{_format_time(self.t_end)}Z"
            f"_v{self.version}{suffix}.fits"
        )

    def relpath(self, label=None):
        """Archive-relative POSIX path, the natural key used by the index.

        Files are filed under the date the observation window *opened*, so a frame
        starting at 23:58 belongs to that day even though it ends on the next one.
        """
        day = datetime.datetime.fromtimestamp(self.t_start, datetime.timezone.utc)
        return (
            f"goes{self.satellite:d}/l2/data/suvi-l2-ci{self.wavelength:03d}/"
            f"{day.year:04d}/{day.month:02d}/{day.day:02d}/{self.filename(label)}"
        )

    def error_plot_name(self):
        """The diagnostic JPG the legacy filter wrote next to a rejected frame."""
        return self.filename()[: -len(".fits")] + "_e.jpg"


def parse_frame_filename(name):
    """Parse a SUVI frame basename, or return None if it is not one.

    Accepts a bare basename only.  Anything carrying a path separator is rejected
    rather than silently normalised, so a crafted name cannot escape the archive
    root when the result is fed back into ``relpath``.
    """
    if not isinstance(name, str) or not name or len(name) > 255:
        return None
    if "/" in name or "\\" in name or os.sep in name:
        return None

    match = _FRAME_RE.match(name)
    if match is None:
        return None

    wavelength = int(match.group("wavelength"))
    satellite = int(match.group("satellite"))
    if wavelength not in WAVELENGTHS or satellite not in SATELLITES:
        return None

    try:
        t_start = _parse_time(match.group("t_start"))
        t_end = _parse_time(match.group("t_end"))
    except ValueError:
        # Syntactically well-formed but not a real date, e.g. month 13.
        return None
    if t_end <= t_start:
        return None

    labels = match.group("labels")
    label = labels[-1] if labels else None

    return FrameName(
        satellite=satellite,
        wavelength=wavelength,
        t_start=t_start,
        t_end=t_end,
        version=match.group("version"),
        label=label,
    )


def data_root():
    """Root of the FITS archive (the ``Data`` directory)."""
    override = os.environ.get("SUVI_DATA_ROOT")
    if override:
        return os.path.abspath(override)
    return os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "Data"))


def default_db_path():
    """Default location of the SQLite index.

    Overridable so the bench can run against a throwaway database without ever
    touching the production index.
    """
    override = os.environ.get("SUVI_DB")
    if override:
        return os.path.abspath(override)
    return os.path.join(os.path.dirname(data_root()), "suvi_index.sqlite")


def abspath(relpath, root=None):
    """Resolve an archive-relative POSIX path against the archive root.

    Raises ValueError if the result would fall outside the root, which keeps a
    malformed or hostile index row from reaching an arbitrary part of the disk.
    """
    root = os.path.abspath(root or data_root())
    resolved = os.path.abspath(os.path.join(root, *relpath.split("/")))
    if resolved != root and not resolved.startswith(root + os.sep):
        raise ValueError(f"Path escapes archive root: {relpath!r}")
    return resolved


def wavelength_dirs(root, satellites=SATELLITES, wavelengths=WAVELENGTHS):
    """Every ``suvi-l2-ciNNN`` directory that could hold frames, existing or not."""
    return [
        os.path.join(root, f"goes{sat:d}", "l2", "data", f"suvi-l2-ci{wl:03d}")
        for sat in satellites
        for wl in wavelengths
    ]