From c3b4e9f3936a357aa8332c230dbcf0e5f1e991b1 Mon Sep 17 00:00:00 2001
From: Sven Sager <akira@narux.de>
Date: Thu, 7 Nov 2024 11:32:34 +0100
Subject: [PATCH] feat: Add context manager for a single device

With this context manager, all IOs can be synchronized from a single
device with the process image.
---
 src/revpimodio2/device.py | 42 +++++++++++++++++++++++++++++++++++++++
 src/revpimodio2/io.py     |  9 ++++++---
 2 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/src/revpimodio2/device.py b/src/revpimodio2/device.py
index 92ce3b6..d4139b7 100644
--- a/src/revpimodio2/device.py
+++ b/src/revpimodio2/device.py
@@ -262,6 +262,48 @@ class Device(object):
         else:
             return key in self._modio.io and getattr(self._modio.io, key)._parentdevice == self
 
+    def __enter__(self):
+        """
+        Read/write inputs/outputs on entering/leaving context manager.
+
+        All inputs of this device are read when entering the context manager.
+        Within the context manager, further .readprocimg() or .writeprocimg()
+        calls can be made. When exiting, all outputs of this device will be
+        written into the process image.
+
+        When 'autorefresh=True' is used, all read or write actions in the
+        background are performed automatically.
+        """
+        if not self._modio._context_manager:
+            # If ModIO itself is in a context manager, it sets the _looprunning=True flag itself
+            if self._modio._looprunning:
+                raise RuntimeError(
+                    "can not enter context manager inside mainloop, cycleloop or "
+                    "another context manager"
+                )
+            self._modio._looprunning = True
+
+        self.readprocimg()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Write outputs of this device to process image on leaving."""
+        if self._modio._imgwriter.is_alive():
+            # Reset new data flag to sync with imgwriter
+            self._modio._imgwriter.newdata.clear()
+
+        # Write outputs on devices without autorefresh
+        if not self._modio._monitoring:
+            self.writeprocimg()
+
+        if self._modio._imgwriter.is_alive():
+            # Wait until imgwriter has written outputs
+            self._modio._imgwriter.newdata.wait(2.5)
+
+        if not self._modio._context_manager:
+            # Do not reset if ModIO is in a context manager itself, it will handle that flag
+            self._modio._looprunning = False
+
     def __getitem__(self, key):
         """
         Gibt IO an angegebener Stelle zurueck.
diff --git a/src/revpimodio2/io.py b/src/revpimodio2/io.py
index 5aa46b8..3c4e061 100644
--- a/src/revpimodio2/io.py
+++ b/src/revpimodio2/io.py
@@ -93,8 +93,8 @@ class IOList(object):
 
         All entries are read when entering the context manager. Within the
         context manager, further .readprocimg() or .writeprocimg() calls can
-        be made and the process image can be read or written. When exiting,
-        all outputs are always written into the process image.
+        be made. When exiting, all outputs will be written into the process
+        image.
 
         When 'autorefresh=True' is used, all read or write actions in the
         background are performed automatically.
@@ -102,7 +102,10 @@ class IOList(object):
         if not self.__modio._context_manager:
             # If ModIO itself is in a context manager, it sets the _looprunning=True flag itself
             if self.__modio._looprunning:
-                raise RuntimeError("can not enter context manager inside mainloop or cycleloop")
+                raise RuntimeError(
+                    "can not enter context manager inside mainloop, cycleloop or "
+                    "another context manager"
+                )
             self.__modio._looprunning = True
 
         self.__modio.readprocimg()