m-labs · Mar 2, 2016 · Mar 2, 2016
Showing with 50 additions and 2 deletions.

+36 −2 artiq/frontend/artiq_rpctool.py

+14 −0 artiq/protocols/pc_rpc.py
diff --git a/artiq/frontend/artiq_rpctool.py b/artiq/frontend/artiq_rpctool.py
@@ -3,10 +3,12 @@
 import argparse
 import textwrap
 import sys
+import traceback
 import numpy as np  # Needed to use numpy in RPC call arguments on cmd line
+import readline  # This makes input() nicer
 import pprint
 
-from artiq.protocols.pc_rpc import AutoTarget, Client
+from artiq.protocols.pc_rpc import AutoTarget, Client, RemoteError
 
 
 def get_argparser():
@@ -17,7 +19,6 @@ def get_argparser():
     parser.add_argument("port", type=int,
                         help="TCP port to use to connect to the controller")
     subparsers = parser.add_subparsers(dest="action")
-    subparsers.required = True
     subparsers.add_parser("list-targets", help="list existing targets")
     parser_list_methods = subparsers.add_parser("list-methods",
                                                 help="list target's methods")
@@ -27,6 +28,10 @@ def get_argparser():
     parser_call.add_argument("method", help="method name")
     parser_call.add_argument("args", nargs=argparse.REMAINDER,
                              help="arguments")
+    parser_interactive = subparsers.add_parser("interactive",
+                                               help="enter interactive mode "
+                                                    "(default)")
+    parser_interactive.add_argument("-t", "--target", help="target name")
     return parser
 
 
@@ -81,8 +86,35 @@ def call_method(remote, method_name, args):
         pprint.pprint(ret)
 
 
+def interactive(remote):
+    while True:
+        try:
+            cmd = input("({}) ".format(remote.get_selected_target()))
+        except EOFError:
+            return
+        class RemoteDict:
+            def __getitem__(self, k):
+                if k == "np":
+                    return np
+                else:
+                    return getattr(remote, k)
+        try:
+            result = eval(cmd, {}, RemoteDict())
+        except Exception as e:
+            if isinstance(e, RemoteError):
+                print("Remote exception:")
+                print(str(e))
+            else:
+                traceback.print_exc()
+        else:
+            if result is not None:
+                print(result)
+
+
 def main():
     args = get_argparser().parse_args()
+    if not args.action:
+        args.target = None
 
     remote = Client(args.server, args.port, None)
     targets, description = remote.get_rpc_id()
@@ -98,6 +130,8 @@ def main():
         list_methods(remote)
     elif args.action == "call":
         call_method(remote, args.method, args.args)
+    elif args.action == "interactive" or not args.action:
+        interactive(remote)
     else:
         print("Unrecognized action: {}".format(args.action))
 

diff --git a/artiq/protocols/pc_rpc.py b/artiq/protocols/pc_rpc.py
@@ -108,6 +108,7 @@ def __init__(self, host, port, target_name=AutoTarget, timeout=None):
             server_identification = self.__recv()
             self.__target_names = server_identification["targets"]
             self.__description = server_identification["description"]
+            self.__selected_target = None
             if target_name is not None:
                 self.select_rpc_target(target_name)
         except:
@@ -119,6 +120,12 @@ def select_rpc_target(self, target_name):
         exactly once if the object was created with ``target_name=None``."""
         target_name = _validate_target_name(target_name, self.__target_names)
         self.__socket.sendall((target_name + "\n").encode())
+        self.__selected_target = target_name
+
+    def get_selected_target(self):
+        """Returns the selected target, or ``None`` if no target has been
+        selected yet."""
+        return self.__selected_target
 
     def get_rpc_id(self):
         """Returns a tuple (target_names, description) containing the
@@ -197,6 +204,7 @@ async def connect_rpc(self, host, port, target_name):
             server_identification = await self.__recv()
             self.__target_names = server_identification["targets"]
             self.__description = server_identification["description"]
+            self.__selected_target = None
             if target_name is not None:
                 self.select_rpc_target(target_name)
         except:
@@ -209,6 +217,12 @@ def select_rpc_target(self, target_name):
         """
         target_name = _validate_target_name(target_name, self.__target_names)
         self.__writer.write((target_name + "\n").encode())
+        self.__selected_target = target_name
+
+    def get_selected_target(self):
+        """Returns the selected target, or ``None`` if no target has been
+        selected yet."""
+        return self.__selected_target
 
     def get_rpc_id(self):
         """Returns a tuple (target_names, description) containing the