From bb4c777ec68c1878bfd132d5268b4b3a686ef147 Mon Sep 17 00:00:00 2001
From: novirium <t.wilson.au@gmail.com>
Date: Tue, 17 Sep 2019 09:34:27 +0800
Subject: [PATCH] TDW update, add convenience methods and allow single args

---
 shepherd/plugins/scout/tdw.py | 40 ++++++++++++++++++++++++++++++++++-
 1 file changed, 39 insertions(+), 1 deletion(-)

diff --git a/shepherd/plugins/scout/tdw.py b/shepherd/plugins/scout/tdw.py
index d8b7628..176a0b6 100644
--- a/shepherd/plugins/scout/tdw.py
+++ b/shepherd/plugins/scout/tdw.py
@@ -47,11 +47,19 @@ MAX_MSG_ARGS = 8
 # dispatch callbacks in new threads asynchronously (or have individual flags when attaching callbacks
 # perhaps?)
 
+class TDWException(Exception):
+    pass
+
+
+class ResponseNotReceivedError(TDWException):
+    pass
+
 
 class MessageType(Enum):
     COMMENT = "#"
     COMMAND = "!"
     REQUEST = "?"
+
     def __str__(self):
         return str(self.value)
 
@@ -86,6 +94,9 @@ class Message():
     def multipart_args(self):
         return self._multipart_args
 
+    def __str__(self):
+        return F"<{self.__class__.__name__}: {self.__dict__}>"
+
 
 class TXMessage(Message):
     """
@@ -95,11 +106,19 @@ class TXMessage(Message):
 
     def __init__(self, msg_type, msg_name, arguments, needs_response, response_timeout, multipart_args=None):
 
+        # Allow single value or string as arg, convert to list
+        if isinstance(arguments, str) or (not hasattr(arguments, "__iter__")):
+            arguments = list(arguments)
+
         # Stringify args and put in immutable tuples
         immutable_args = tuple([str(arg) for arg in arguments])
 
         immutable_multipart_args = None
         if multipart_args is not None:
+            # Allow single value or string as arg, convert to list
+            if isinstance(multipart_args, str) or (not hasattr(multipart_args, "__iter__")):
+                multipart_args = list(multipart_args)
+
             immutable_multipart_args = []
             for arglist in multipart_args:
                 immutable_multipart_args.append(tuple([str(arg) for arg in arglist]))
@@ -236,12 +255,31 @@ class MessageHandler():
         self.send_message(msg)
         return msg
 
+    def response_from_request(self, message_name, arguments=[], response_timeout=DEFAULT):
+        '''
+        Sends request and returns the response. Blocks while waiting. Throws ResponseNotReceivedError
+        if the response times out.
+        '''
+        rqst = self.send_request(message_name, arguments, True, response_timeout)
+        if rqst.wait_for_response():
+            return rqst.response.arguments
+        raise ResponseNotReceivedError(rqst)
+
+    def response_from_command(self, message_name, arguments=[], response_timeout=DEFAULT):
+        '''
+        Sends command and returns the response. Blocks while waiting. Throws ResponseNotReceivedError
+        if the response times out.
+        '''
+        cmd = self.send_command(message_name, arguments, True, response_timeout)
+        if cmd.wait_for_response():
+            return cmd.response.arguments
+        raise ResponseNotReceivedError(cmd)
+
     def _send_message(self):
         """
         Actually send a message pulled from the queue.
         Only called from the serial_comm thread
         """
-        
 
         argstr = ""
         if len(self._tx_message.arguments) > 0: