use re-implemented HDUList

csst_common/fits.py

@@ -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(
-            filename, use_fsspec=True, fsspec_kwargs=s3_options, **kwargs
+        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