Allow to use fitsio to save files (#45)

* Ignore astropy typings

* Do not use user config file

* Support using fitsio

* Fix unnecessary import

* Update CHANGELOG.md

* Decreate codecov target for project

* Test invalid write engine

* Add test for post_process

* Test import error without fitsio

.github/workflows/test.yml

@@ -35,7 +35,7 @@ jobs:
       - name: Install dependencies
         run: |
           pip install --upgrade wheel pip setuptools
-          pip install .
+          pip install .[fitsio]
 
       - name: Lint with ruff
         run: |

CHANGELOG.md

@@ -1,5 +1,12 @@
 # Changelog
 
+## Next version
+
+### 🚀 New
+
+* [#45](https://github.com/sdss/archon/issues/45) Added a new option `files.write_engine` that can be set to `astropy` or `fitsio`. In the latter case it will use fitsio to write images to disk. This requires installing `sdss-archon` with the `fitsio` extra (e.g., `pip install sdss-archon[fitsio]`).
+
+
 ## 0.11.6 - November 24, 2023
 
 ### 🏷️ Changed

archon/__init__.py

@@ -7,7 +7,7 @@
 NAME = "sdss-archon"
 
 # Loads config. config name is the package name.
-config = get_config("archon")
+config = get_config("archon", allow_user=False)
 
 # Inits the logging system as NAME. Remove all the handlers. If a client
 # of the library wants the archon logging, it should add its own handler.

archon/actor/delegate.py

@@ -20,12 +20,13 @@
 from time import time
 from unittest.mock import MagicMock
 
-from typing import TYPE_CHECKING, Any, Dict, Generic, List, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Dict, Generic, List, Sequence, TypeVar, cast
 
 import astropy.time
 import numpy
 from astropy.io import fits
 
+from sdsstools.configuration import Configuration
 from sdsstools.time import get_sjd
 
 from archon import __version__
@@ -59,6 +60,8 @@ def __init__(self, actor: Actor_co):
         self._command: Command[Actor_co] | None = None
         self._expose_cotasks: asyncio.Task | None = None
 
+        self._check_fitsio()
+
     @property
     def command(self):
         """Returns the current command."""
@@ -298,19 +301,20 @@ async def readout(
             self.reset()
             return True
 
-        c_to_hdus: dict[ArchonController, list[fits.PrimaryHDU]]
+        c_to_hdus: dict[ArchonController, list[dict[str, Any]]]
         c_to_hdus = {controllers[ii]: hdus[ii] for ii in range(len(controllers))}
 
         post_process_jobs = []
         for controller, hdus in c_to_hdus.items():
             post_process_jobs.append(self.post_process(controller, hdus))
-
         c_to_hdus = dict(list(await asyncio.gather(*post_process_jobs)))
+
         self.command.debug(text="Writing HDUs to file.")
-        await asyncio.gather(*[self.write_hdus(c, h) for c, h in c_to_hdus.items()])
+        coros = [self.write_hdus(c, h) for c, h in c_to_hdus.items()]
+        results = await asyncio.gather(*coros)
 
         self.reset()
-        return True
+        return all(results)
 
     async def expose_cotasks(self):
         """Tasks that will be executed concurrently with readout.
@@ -353,25 +357,25 @@ async def fetch_hdus(self, controller: ArchonController) -> List[fits.PrimaryHDU
         self.expose_data.header["BUFFER"] = (buffer_no, "The buffer number read")
 
         controller_info = config["controllers"][controller.name]
-        hdus = []
+        hdus: list[dict[str, Any]] = []
         for ccd_name in controller_info["detectors"]:
             header = await self.build_base_header(controller, ccd_name)
             ccd_data = self._get_ccd_data(data, controller, ccd_name, controller_info)
-            hdus.append(fits.PrimaryHDU(data=ccd_data, header=header))
+            hdus.append({"data": ccd_data, "header": header})
 
         return hdus
 
     async def write_hdus(
         self,
         controller: ArchonController,
-        hdus: List[fits.PrimaryHDU],
+        hdus: list[dict[str, Any]],
     ):
         """Writes HDUs to disk."""
 
         assert self.expose_data
 
         expose_data = self.expose_data
-        config = self.actor.config
+        config = Configuration(self.actor.config)
 
         excluded_cameras = config.get("excluded_cameras", [])
 
@@ -380,20 +384,16 @@ async def write_hdus(
 
         write_tasks = []
         for _, hdu in enumerate(hdus):
-            ccd = cast(str, hdu.header["ccd"])
+            ccd = cast(str, hdu["header"]["CCD"][0])
 
             if ccd in excluded_cameras:
                 self.command.warning(f"Not saving image for camera {ccd}.")
                 continue
 
             file_path = self._get_ccd_filepath(controller, ccd)
 
-            hdu.header["filename"] = os.path.basename(file_path)
-            hdu.header.insert(
-                "filename",
-                ("EXPOSURE", expose_data.exposure_no, "Exposure number"),
-                after=True,
-            )
+            hdu["header"]["FILENAME"][0] = os.path.basename(file_path)
+            hdu["header"]["EXPOSURE"][0] = expose_data.exposure_no
 
             write_tasks.append(
                 self._write_to_file(
@@ -432,11 +432,11 @@ async def write_hdus(
 
     async def _write_to_file(
         self,
-        hdu: fits.PrimaryHDU,
+        data: dict[str, Any],
         file_path: str,
         write_async: bool = True,
     ):
-        """Writes the HDU to file using an executor.
+        """Writes the HDU to file using astropy or fitsio.
 
         The file is first written to a temporary file with the same path and
         name as the final file but with a random suffix, and then renamed.
@@ -446,10 +446,17 @@ async def _write_to_file(
         if os.path.exists(file_path):
             raise FileExistsError(f"Cannot overwrite file {file_path}.")
 
-        loop = asyncio.get_running_loop()
+        loop = asyncio.get_event_loop()
+
+        write_engine: str = self.actor.config["files"].get("write_engine", "astropy")
+        if write_engine == "astropy":
+            writeto = partial(self._write_file_astropy, data)
+        elif write_engine == "fitsio":
+            writeto = partial(self._write_file_fitsio, data)
+        else:
+            raise ValueError(f"Invalid write engine {write_engine!r}.")
 
-        writeto = partial(hdu.writeto, checksum=True)
-        temp_file = NamedTemporaryFile(suffix=".fits", delete=False).name
+        temp_file = NamedTemporaryFile(suffix=".fits", delete=True).name
 
         if write_async:
             if file_path.endswith(".gz"):
@@ -485,6 +492,36 @@ async def _write_to_file(
 
         return file_path
 
+    def _write_file_astropy(self, data: dict[str, Any], file_path: str):
+        """Writes the HDU to file using astropy."""
+
+        header = fits.Header()
+        for key, value in data["header"].items():
+            header[key] = tuple(value) if isinstance(value, (list, tuple)) else value
+
+        hdu = fits.PrimaryHDU(data["data"], header=header)
+        hdu.writeto(file_path, checksum=True, overwrite=True)
+
+        return
+
+    def _write_file_fitsio(self, data: dict[str, Any], file_path: str):
+        """Writes the HDU to file using astropy."""
+
+        import fitsio
+
+        header = []
+        for key, value in data["header"].items():
+            if isinstance(value, Sequence):
+                header.append({"name": key, "value": value[0], "comment": value[1]})
+            else:
+                header.append({"name": key, "value": value, "comment": ""})
+
+        with fitsio.FITS(file_path, "rw") as fits_:
+            fits_.write(data["data"], header=header)
+            fits_[-1].write_checksum()
+
+        return
+
     async def _generate_checksum(self, filenames: list[str]):
         """Generates a checksum file for the images written to disk."""
 
@@ -527,8 +564,8 @@ async def _generate_checksum(self, filenames: list[str]):
     async def post_process(
         self,
         controller: ArchonController,
-        hdus: list[fits.PrimaryHDU],
-    ) -> tuple[ArchonController, list[fits.PrimaryHDU]]:
+        hdus: list[dict[str, Any]],
+    ) -> tuple[ArchonController, list[dict[str, Any]]]:
         """Custom post-processing."""
 
         return (controller, hdus)
@@ -541,11 +578,12 @@ async def build_base_header(self, controller: ArchonController, ccd_name: str):
         expose_data = self.expose_data
         assert expose_data.end_time is not None
 
-        header = fits.Header()
+        header: dict[str, tuple[Any, str] | Any] = {}
 
         # Basic header
         header["V_ARCHON"] = __version__
-        header["FILENAME"] = ("", "File basename")  # Will be filled out later
+        header["FILENAME"] = ["", "File basename"]  # Will be filled out later
+        header["EXPOSURE"] = [None, "Exposure number"]  # Will be filled out later
         header["SPEC"] = (controller.name, "Spectrograph name")
         header["OBSERVAT"] = (self.command.actor.observatory, "Observatory")
         header["OBSTIME"] = (expose_data.start_time.isot, "Start of the observation")
@@ -558,15 +596,18 @@ async def build_base_header(self, controller: ArchonController, ccd_name: str):
 
         header["CCD"] = (ccd_name, "CCD name")
 
-        config = self.actor.config
-        if (
-            "controllers" not in config or controller.name not in config["controllers"]
-        ):  # pragma: no cover
+        config = Configuration(self.actor.config)
+        controller_config = config[f"controllers.{controller.name}"]
+        if controller_config is None:  # pragma: no cover
             self.command.warning(text="Cannot retrieve controller information.")
-            controller_config = {"detectors": {ccd_name: {}}, "parameters": {}}
-        else:
-            controller_config = config["controllers"][controller.name]
-        ccd_config = controller_config["detectors"][ccd_name]
+            controller_config = Configuration(
+                {
+                    "detectors": {ccd_name: {}},
+                    "parameters": {},
+                }
+            )
+
+        ccd_config = controller_config[f"detectors.{ccd_name}"]
 
         ccdid = ccd_config.get("serial", "?")
         ccdtype = ccd_config.get("type", "?")
@@ -846,6 +887,19 @@ def _get_ccd_data(
 
         return ccd_data
 
+    def _check_fitsio(self):
+        """Checks if fitsio is installed and needed."""
+
+        write_engine: str = self.actor.config["files"].get("write_engine", "astropy")
+        if write_engine == "fitsio":
+            try:
+                import fitsio  # noqa: F401
+            except ImportError:
+                raise ImportError(
+                    "fitsio is required to use fitsio. You can install "
+                    "it with 'pip install fitsio' or 'pip install sdss-archon[fitsio]'."
+                )
+
 
 @dataclass
 class ExposeData:

archon/etc/archon.yml

@@ -17,7 +17,8 @@ timeouts:
 files:
   data_dir: '~/'
   template: 'sdR-{ccd}-{exposure_no:08d}.fits.gz'
-  write_async: false
+  write_async: true
+  write_engine: astropy
 
 archon:
   default_parameters: {}

codecov.yml

@@ -5,7 +5,7 @@ coverage:
   status:
     project:
       default:
-        target: 90%
+        target: 80%
         if_not_found: success
         if_ci_failed: error
         informational: false

pyproject.toml

@@ -25,11 +25,12 @@ packages = [
 
 [tool.poetry.dependencies]
 python = "^3.9,<3.13"
-sdsstools = "^1.0.2"
+sdsstools = "^1.5.2"
 numpy = "^1.26.1"
 sdss-clu = "^2.0.0b2"
 click-default-group = "^1.2.2"
 astropy = "^6.0"
+fitsio = {version = "^1.2.1", optional = true}
 
 [tool.poetry.group.dev.dependencies]
 ipython = [
@@ -56,6 +57,9 @@ sphinx-autobuild = ">=2021.3.14"
 sphinx-copybutton = ">=0.3.3"
 ruff = ">=0.0.284"
 
+[tool.poetry.extras]
+fitsio = ["fitsio"]
+
 [tool.black]
 line-length = 88
 target-version = ['py311']
@@ -65,6 +69,7 @@ fast = true
 line-length = 88
 target-version = 'py311'
 select = ["E", "F", "I"]
+exclude = ["__init__.pyi"]
 
 [tool.ruff.per-file-ignores]
 "__init__.py" = ["F403", "F401", "E402"]