From 359123b22bd2c1bda8d14c8d6633047a30a4b1a2 Mon Sep 17 00:00:00 2001
From: Lee Begg <lee@begg.nz>
Date: Sat, 22 May 2021 21:11:37 +1200
Subject: [PATCH] Expose the mqttc loop functions, and option to avoid the auto
 start

This also makes a cleaner `main()` by using `loop_forever()`

Default for auto_start_loop is backwards compatible and structured to
allow for changing the default later and/or raising a warning if not
set.
---
 sondehub/__init__.py | 30 ++++++++++++++++++++++++++++--
 sondehub/__main__.py |  5 ++---
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/sondehub/__init__.py b/sondehub/__init__.py
index ea4cf7a..88f5263 100644
--- a/sondehub/__init__.py
+++ b/sondehub/__init__.py
@@ -13,6 +13,8 @@ import queue
 
 S3_BUCKET = "sondehub-open-data"
 
+class LoopStartedError(Exception):
+    pass
 
 class Stream:
     def __init__(self,
@@ -20,7 +22,8 @@ class Stream:
                  on_connect=None,
                  on_message=None,
                  on_log=None,
-                 on_disconnect=None, asJson=False):
+                 on_disconnect=None, asJson=False,
+                 auto_start_loop=None):
         self.mqttc = mqtt.Client(transport="websockets")
         self._sondes = sondes
         self.asJson = asJson
@@ -28,6 +31,9 @@ class Stream:
         self.on_message = on_message
         self.on_disconnect = on_disconnect
         self.on_log = on_log
+        if auto_start_loop is None:
+            auto_start_loop = True
+        self.auto_start_loop = auto_start_loop
         self.ws_connect()
 
 
@@ -66,7 +72,8 @@ class Stream:
             self.mqttc.connect(urlparts.netloc, 443, 60)
         except OSError:
             pass
-        self.mqttc.loop_start()
+        if self.auto_start_loop:
+            self.mqttc.loop_start()
 
     def get_url(self):
         conn = http.client.HTTPSConnection("api.v2.sondehub.org")
@@ -108,6 +115,25 @@ class Stream:
 
     def disconnect(self):
         self.mqttc.disconnect()
+    
+    def loop_start(self):
+        if self.auto_start_loop:
+            raise LoopStartedError()
+        self.mqttc.loop_start()
+    
+    def loop_stop(self):
+        self.mqttc.loop_stop()
+        self.auto_start_loop = False
+    
+    def loop_step(self, timeout=1.0):
+        if self.auto_start_loop:
+            raise LoopStartedError()        
+        self.mqttc.loop(timeout)
+    
+    def loop_forever(self):
+        if self.auto_start_loop:
+            raise LoopStartedError()        
+        self.mqttc.loop_forever()
 
 
 class Downloader(threading.Thread):
diff --git a/sondehub/__main__.py b/sondehub/__main__.py
index e89e1aa..366241a 100644
--- a/sondehub/__main__.py
+++ b/sondehub/__main__.py
@@ -47,9 +47,8 @@ def main():
     ):  # we need to drop the default value if the user specifies specific sondes
         args.sondes = args.sondes[1:]
     sondes = [item for sublist in args.sondes for item in sublist]
-    test = sondehub.Stream(on_message=on_message, sondes=sondes)
-    while 1:
-        time.sleep(0.01) # don't overwork the CPU waiting for events
+    test = sondehub.Stream(on_message=on_message, sondes=sondes, auto_start_loop=False)
+    test.loop_forever()
 
 
 if __name__ == "__main__":