Commit 36ed8333 authored by BO ZHANG's avatar BO ZHANG 🏀
Browse files

add HDUList.to_bytes()

parent b6e173f9

csst_common/fits.py

+27 −13
Original line number Diff line number Diff line
@@ -8,8 +8,9 @@ Modified-History: 
    2025-10-23, Bo Zhang, created, support FITS operations for s3

"""



import io

from astropy.io.fits import *

from astropy.io import fits

from astropy.io import fits as astropy_fits

from csst_fs.s3_config import load_s3_options



s3_options = load_s3_options()
@@ -17,43 +18,56 @@ s3_options = load_s3_options() 
s3_prefix = ("s3://", "s3a://")





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

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

    if filename.startswith(s3_prefix):

        # read FITS file from s3

        return fits.open(filename, use_fsspec=True, fsspec_kwargs=s3_options, **kwargs)

        return astropy_fits.open(

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

        )

    else:

        # read FITS file from local

        return fits.open(filename, **kwargs)

        return astropy_fits.open(filename, **kwargs)





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

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

    if filename.startswith(s3_prefix):

        # read FITS file from s3

        return fits.getheader(

        return astropy_fits.getheader(

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

        )

    else:

        # read FITS file from local

        return fits.getheader(filename, **kwargs)

        return astropy_fits.getheader(filename, **kwargs)





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

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

    if filename.startswith(s3_prefix):

        # read FITS file from s3

        return fits.getval(

        return astropy_fits.getval(

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

        )

    else:

        # read FITS file from local

        return fits.getval(filename, **kwargs)

        return astropy_fits.getval(filename, **kwargs)





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

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

    if filename.startswith(s3_prefix):

        # read FITS file from s3

        return fits.getdata(

        return astropy_fits.getdata(

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

        )

    else:

        # read FITS file from local

        return fits.getdata(filename, **kwargs)
 
        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

tests/test_fits.py

+5 −0
Original line number Diff line number Diff line
@@ -47,3 +47,8 @@ class TestFitsHeaderOps(unittest.TestCase): 
        data = fits.getdata(test_fits_file, ext=1)

        self.assertIsInstance(data, np.ndarray)

        self.assertEqual(data.shape, (5, 5))



    def test_fits_to_bytes(self):

        hl = fits.open(test_fits_file)

        fits_bytes = hl.to_bytes()

        self.assertIsInstance(fits_bytes, bytes)