Commit 21e2d3d6 authored by BO ZHANG's avatar BO ZHANG 🏀
Browse files

use re-implemented HDUList

parent f56bb997

csst_common/fits.py

+36 −33
Original line number Diff line number Diff line
@@ -14,24 +14,53 @@ import os.path 
from astropy.io import fits as astropy_fits

from astropy.io.fits import *

from csst_fs.s3_config import load_s3_options

from typing import Any



s3_options = load_s3_options()



s3_prefix = ("s3://", "s3a://")





def open(filename, **kwargs) -> astropy_fits.HDUList:

class HDUList(astropy_fits.HDUList):

    def to_bytes(self) -> bytes:

        """Convert the HDUList to bytes."""



        # 使用 with 语句自动管理资源

        with io.BytesIO() as memory_buffer:

            self.writeto(memory_buffer)

            fits_bytes = memory_buffer.getvalue()

        return fits_bytes





def prepare_fits_files_for_ingestion(files: list[str]) -> list[dict]:

    """Prepare FITS files for ingestion."""

    ingestion_files = []

    for file in files:

        with open(file) as hl:

            fits_bytes = hl.to_bytes()

        ingestion_files.append(

            {

                "file_name": os.path.basename(file),

                "file_content": fits_bytes,

            }

        )

    return ingestion_files





def open(filename, **kwargs) -> HDUList:

    if filename.startswith(s3_prefix):

        # read FITS file from s3

        return astropy_fits.open(

        return HDUList(

            astropy_fits.open(

                filename, use_fsspec=True, fsspec_kwargs=s3_options, **kwargs

            )

        )

    else:

        # read FITS file from local

        return astropy_fits.open(filename, **kwargs)

        return HDUList(astropy_fits.open(filename, **kwargs))





def getheader(filename, **kwargs) -> astropy_fits.HDUList:

def getheader(filename, **kwargs) -> astropy_fits.Header:

    if filename.startswith(s3_prefix):

        # read FITS file from s3

        return astropy_fits.getheader(
@@ -42,7 +71,7 @@ def getheader(filename, **kwargs) -> astropy_fits.HDUList: 
        return astropy_fits.getheader(filename, **kwargs)





def getval(filename, **kwargs) -> astropy_fits.HDUList:

def getval(filename, **kwargs) -> Any:

    if filename.startswith(s3_prefix):

        # read FITS file from s3

        return astropy_fits.getval(
@@ -53,7 +82,7 @@ def getval(filename, **kwargs) -> astropy_fits.HDUList: 
        return astropy_fits.getval(filename, **kwargs)





def getdata(filename, **kwargs) -> astropy_fits.HDUList:

def getdata(filename, **kwargs) -> Any:

    if filename.startswith(s3_prefix):

        # read FITS file from s3

        return astropy_fits.getdata(
@@ -62,29 +91,3 @@ def getdata(filename, **kwargs) -> astropy_fits.HDUList: 
    else:

        # read FITS file from local

        return astropy_fits.getdata(filename, **kwargs)
 




class HDUList(astropy_fits.HDUList):

    def to_bytes(self) -> bytes:

        """Convert the HDUList to bytes."""



        # 使用 with 语句自动管理资源

        with io.BytesIO() as memory_buffer:

            self.writeto(memory_buffer)

            fits_bytes = memory_buffer.getvalue()

        return fits_bytes





def prepare_fits_files_for_ingestion(files: list[str]) -> list[dict]:

    """Prepare FITS files for ingestion."""

    ingestion_files = []

    for file in files:

        with open(file) as hl:

            fits_bytes = hl.to_bytes()

        ingestion_files.append(

            {

                "file_name": os.path.basename(file),

                "file_content": fits_bytes,

            }

        )

    return ingestion_files