Commit be7eda29 authored by BO ZHANG's avatar BO ZHANG 🏀
Browse files

add delete_section

parent fb155307
Loading
Loading
Loading
Loading
+57 −0
Original line number Diff line number Diff line
@@ -8,7 +8,10 @@ Modified-History:
    2023-12-10, Bo Zhang, created
    2023-12-15, Bo Zhang, added verify_checksum and append_header, add module header
    2023-12-16, Bo Zhang, fixed a bug in finding the first key when reformatting
    2024-06-08, Bo Zhang, added delete_section
"""

import bisect
import warnings
from copy import deepcopy
from typing import Optional
@@ -134,3 +137,57 @@ def reformat_header(
    h_copy.add_comment(comment, before=first_key)
    h_copy.add_comment("=" * 72, before=first_key)
    return h_copy


def delete_section(
    header: fits.Header,
    title="WORLD COORDINATE SYSTEM INFORMATION",
) -> fits.Header:
    """
    Delete a section from a fits header.

    This function is used to delete the section of WCS information from the header.

    Parameters
    ----------
    header: fits.Header
        Header.
    title: str
        Title of the section to be deleted.

    Returns
    -------
    fits.Header
        Edited header.
    """

    # find sep lines
    i_seps = []
    for i_card, card in enumerate(header.cards):
        if (
            card.keyword == "COMMENT"
            and isinstance(card.value, str)
            and card.value.startswith("===")
        ) or i_card == len(header.cards) - 1:
            i_seps.append(i_card)

    # find section title
    i_section = -1
    for i_card, card in enumerate(header.cards):
        if card.keyword == "COMMENT" and card.value == title:
            i_section = i_card
            break
    if i_section < 0:
        raise ValueError(f"Title not found: `{title}`")

    # find section start and stop
    idx = bisect.bisect(i_seps, i_section)
    i_section_start = i_seps[idx - 1]
    i_section_stop = i_seps[idx + 1]

    # delete keys in reverse order
    print(f"Delete keys {i_section_start} to {i_section_stop} ...")
    for i in range(i_section_stop - 1, i_section_start - 1, -1):
        del header[i]

    return header