From bb4a7b86b4d4bd2c84bea3459b3768415394ed96 Mon Sep 17 00:00:00 2001
From: Tim-Oliver Husser <thusser@uni-goettingen.de>
Date: Mon, 25 Dec 2023 12:58:19 +0100
Subject: [PATCH] options for not storing focus series but take a final image

---
 pyobs/modules/focus/focusseries.py | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/pyobs/modules/focus/focusseries.py b/pyobs/modules/focus/focusseries.py
index de7c0cf5..8c9a5967 100644
--- a/pyobs/modules/focus/focusseries.py
+++ b/pyobs/modules/focus/focusseries.py
@@ -31,6 +31,8 @@ def __init__(
         filters: Optional[Union[str, IFilters]] = None,
         filter_name: Optional[str] = None,
         binning: Optional[int] = None,
+        broadcast: bool = False,
+        final_image: bool = True,
         **kwargs: Any,
     ):
         """Initialize a new auto focus system.
@@ -51,6 +53,8 @@ def __init__(
         self._offset = offset
         self._abort = threading.Event()
         self._running = False
+        self._broadcast = broadcast
+        self._final_image = final_image
 
         # create focus series
         self._series: FocusSeries = get_object(series, FocusSeries)
@@ -122,6 +126,8 @@ async def _auto_focus(self, count: int, step: float, exposure_time: float, **kwa
 
         # do camera settings
         await self._do_camera_settings(camera)
+        if isinstance(camera, IImageType):
+            await camera.set_image_type(ImageType.FOCUS)
 
         # get filter wheel and current filter
         filter_name = "unknown"
@@ -173,11 +179,7 @@ async def _auto_focus(self, count: int, step: float, exposure_time: float, **kwa
             if self._abort.is_set():
                 raise exceptions.AbortedError()
             try:
-                if isinstance(camera, IExposureTime):
-                    await camera.set_exposure_time(exposure_time)
-                if isinstance(camera, IImageType):
-                    await camera.set_image_type(ImageType.FOCUS)
-                filename = await camera.grab_data()
+                filename = await self._take_image(camera, exposure_time)
             except exc.RemoteError:
                 log.error("Could not take image.")
                 continue
@@ -231,9 +233,21 @@ async def _auto_focus(self, count: int, step: float, exposure_time: float, **kwa
         # send event
         await self.comm.send_event(FocusFoundEvent(absolute, focus[1], filter_name))
 
+        # take final image?
+        if self._final_image:
+            await self._take_image(camera, exposure_time)
+
         # return result
         return focus[0], focus[1]
 
+    async def _take_image(self, camera, exposure_time):
+        if isinstance(camera, IExposureTime):
+            await camera.set_exposure_time(exposure_time)
+        if isinstance(camera, IData):
+            return await camera.grab_data(broadcast=self._broadcast)
+        else:
+            raise exc.GeneralError("Cannot grab data from camera.")
+
     async def auto_focus_status(self, **kwargs: Any) -> Dict[str, Any]:
         """Returns current status of auto focus.