From 49b60edf7214d44fdc81f7464c76fbb08d07aaf7 Mon Sep 17 00:00:00 2001
From: Cameron Devine <cameron.devine@alleninstitute.org>
Date: Tue, 25 Jun 2024 17:04:10 -0700
Subject: [PATCH] added an options to the CLI to subscribe to all topics, and
 to exit after one message is received.

---
 pigeon/__main__.py | 42 ++++++++++++++++++++++++++++--------------
 pyproject.toml     |  4 ++--
 2 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/pigeon/__main__.py b/pigeon/__main__.py
index a65cb1b..3315df9 100644
--- a/pigeon/__main__.py
+++ b/pigeon/__main__.py
@@ -3,9 +3,14 @@
 import yaml
 
 
-def callback(topic, msg):
-    print(f"Recieved message on topic '{topic}':")
-    print(msg)
+class Listener:
+    def __init__(self):
+        self.message_received = False
+
+    def callback(self, topic, msg):
+        print(f"Recieved message on topic '{topic}':")
+        print(msg)
+        self.message_received = True
 
 
 def main():
@@ -17,6 +22,8 @@ def main():
     parser.add_argument("-p", "--publish", type=str, help="The topic to publish a message to.")
     parser.add_argument("-d", "--data", type=str, help="The YAML/JSON formatted data to publish.")
     parser.add_argument("-s", "--subscribe", type=str, action="append", default=[], help="The topic to subscribe to.")
+    parser.add_argument("--one", action="store_true", help="Exit after receiving one message.")
+    parser.add_argument("-a", "--all", action="store_true", help="Subscribe to all registered topics.")
     parser.add_argument("-l", "--list", action="store_true", help="List registered topics and exit.")
 
     args = parser.parse_args()
@@ -27,34 +34,41 @@ def main():
             print(topic)
         return
 
-    if args.publish is None and args.subscribe is None:
+    if args.publish is None and args.subscribe is None and not args.all:
         print("No action specified.")
         return
-    
+
     if args.publish and args.data is None:
         print("Must also specify data to publish.")
         return
-    
+
     if args.data and args.publish is None:
         print("Most also specify topic to publish data to.")
         return
-    
+
     connection = Pigeon("CLI", args.host, args.port)
     connection.connect(args.username, args.password)
-    
+
     if args.publish:
         connection.send(args.publish, **yaml.safe_load(args.data))
-    
-    for topic in args.subscribe:
-        connection.subscribe(topic, callback)
 
-    if args.subscribe:
+    if args.subscribe or args.all:
+        listener = Listener()
+
+    if args.all:
+        connection.subscribe_all(listener.callback)
+    else:
+        for topic in args.subscribe:
+            connection.subscribe(topic, listener.callback)
+
+    if args.subscribe or args.all:
         try:
-            while True:
+            while not (args.one and listener.message_received):
                 pass
         except KeyboardInterrupt:
             print("exiting")
+    exit(0)
 
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()
diff --git a/pyproject.toml b/pyproject.toml
index d59b9e2..36eee5e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "pigeon-client"
-version = "0.1.1"
+version = "0.1.2"
 authors = [
   { name="Cameron Devine", email="cameron.devine@alleninstitute.org" },
 ]
@@ -27,4 +27,4 @@ Repository = "https://github.com/AllenInstitute/pigeon"
 Issues = "https://github.com/AllenInstitute/pigeon/issues"
 
 [tool.setuptools.dynamic]
-dependencies = {file = ["requirements.txt"]}
\ No newline at end of file
+dependencies = {file = ["requirements.txt"]}