on_close

Author: joente

README.md

@@ -115,6 +115,19 @@ if __name__ == "__main__":
     probe.start()
 ```
 
+## On CLose
+
+Using the `set_on_close()` method, it is possible to configure a method which will be called before the probe is stopped.
+This can be useful in case you want to nicely close some connections.
+
+```python
+async def custom_on_close():
+    ...
+
+probe = Probe("myProbe", '3.0.0', {})
+probe.set_on_close(custom_on_close)
+```
+
 ## ASCII item names
 
 InfraSonar requires each item to have a unique _name_ property. The value for _name_ must be a _string_ with ASCII compatible character.

libprobe/probe.py

@@ -136,6 +136,7 @@ def __init__(
         self._checks: Dict[Tuple[int, int], asyncio.Future] = {}
         self._dry_run: Optional[Tuple[Asset, dict]] = \
             None if dry_run is None else self._load_dry_run_assst(dry_run)
+        self._on_close: Callable[[], Awaitable[None]] | None = None
 
         if not os.path.exists(config_path):
             try:
@@ -202,11 +203,17 @@ def is_connected(self) -> bool:
     def is_connecting(self) -> bool:
         return self._connecting
 
+    def set_on_close(self, on_close: Callable[[], Awaitable[None]]):
+        self._on_close = on_close
+
     def _stop(self, signame, *args):
         logging.warning(
             f'signal \'{signame}\' received, stop {self.name} probe')
-        for task in asyncio.all_tasks():
-            task.cancel()
+        if self._on_close is not None and self.loop is not None:
+            self.loop.run_until_complete(self._on_close())
+        else:
+            for task in asyncio.all_tasks():
+                task.cancel()
 
     async def _start(self):
         initial_step = 2
@@ -335,6 +342,8 @@ async def _do_dry_run(self):
             output = json.dumps(response, indent=2)
             print(output)
         print('', file=sys.stderr)
+        if self._on_close is not None:
+            await self._on_close()
 
     async def _connect(self):
         assert self.loop is not None

libprobe/version.py

@@ -1 +1 @@
-__version__ = '1.0.1'
+__version__ = '1.0.2'