add HDUList.to_bytes()

csst_common/fits.py

@@ -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

@@ -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)