m-labs · Oct 25, 2014 · Oct 25, 2014 · Oct 25, 2014 · Oct 25, 2014
Showing with 83 additions and 8 deletions.

+1 −1 artiq/management/pc_rpc.py

+8 −0 artiq/management/pyon.py

+5 −4 artiq/management/worker.py

+4 −3 artiq/management/worker_impl.py

+65 −0 test/pc_rpc.py
diff --git a/artiq/management/pc_rpc.py b/artiq/management/pc_rpc.py
@@ -77,8 +77,8 @@ def handle_connection_task(self, reader, writer):
                 obj = pyon.decode(line.decode())
                 action = obj["action"]
                 if action == "call":
-                    method = getattr(self.target, obj["name"])
                     try:
+                        method = getattr(self.target, obj["name"])
                         ret = method(*obj["args"], **obj["kwargs"])
                         obj = {"result": "ok", "ret": ret}
                     except Exception as e:

diff --git a/artiq/management/pyon.py b/artiq/management/pyon.py
@@ -7,6 +7,13 @@ def _encode_none(x):
     return "None"
 
 
+def _encode_bool(x):
+    if x:
+        return "True"
+    else:
+        return "False"
+
+
 def _encode_number(x):
     return str(x)
 
@@ -50,6 +57,7 @@ def _encode_nparray(x):
 
 _encode_map = {
     type(None): _encode_none,
+    bool: _encode_bool,
     int: _encode_number,
     float: _encode_number,
     str: _encode_str,

diff --git a/artiq/management/worker.py b/artiq/management/worker.py
@@ -2,7 +2,8 @@
 import asyncio
 import subprocess
 import signal
-import json
+
+from artiq.management import pyon
 
 
 class WorkerFailed(Exception):
@@ -24,7 +25,7 @@ def create_process(self):
 
     @asyncio.coroutine
     def _send(self, obj, timeout):
-        line = json.dumps(obj)
+        line = pyon.encode(obj)
         self.process.stdin.write(line.encode())
         self.process.stdin.write("\n".encode())
         try:
@@ -47,9 +48,9 @@ def _recv(self, timeout):
             raise WorkerFailed(
                 "Worker ended unexpectedly while trying to receive data")
         try:
-            obj = json.loads(line.decode())
+            obj = pyon.decode(line.decode())
         except:
-            raise WorkerFailed("Worker sent invalid JSON data")
+            raise WorkerFailed("Worker sent invalid PYON data")
         return obj
 
     @asyncio.coroutine

diff --git a/artiq/management/worker_impl.py b/artiq/management/worker_impl.py
@@ -1,7 +1,8 @@
-import json
 import sys
 import importlib
 
+from artiq.management import pyon
+
 
 def import_in_folder(path, name):
     try:
@@ -20,7 +21,7 @@ def run(path, name):
 
 
 def put_object(obj):
-    ds = json.dumps(obj)
+    ds = pyon.encode(obj)
     sys.__stdout__.write(ds)
     sys.__stdout__.write("\n")
     sys.__stdout__.flush()
@@ -31,7 +32,7 @@ def main():
 
     while True:
         line = sys.__stdin__.readline()
-        obj = json.loads(line)
+        obj = pyon.decode(line)
         put_object("ack")
 
         try:

diff --git a/test/pc_rpc.py b/test/pc_rpc.py
@@ -0,0 +1,65 @@
+import unittest
+import sys
+import subprocess
+import asyncio
+import time
+
+import numpy as np
+
+from artiq.management import pc_rpc
+
+
+test_address = "::1"
+test_port = 7777
+
+
+class RPCCase(unittest.TestCase):
+    def test_echo(self):
+        # running this file outside of unittest starts the echo server
+        with subprocess.Popen([sys.executable,
+                               sys.modules[__name__].__file__]) as proc:
+            try:
+                test_object = [5, 2.1, None, True, False,
+                               {"a": 5, 2: np.linspace(0, 10, 1)},
+                               (4, 5), (10,)]
+                time.sleep(.5)  # wait for the server to start
+                remote = pc_rpc.Client(test_address, test_port)
+                try:
+                    test_object_back = remote.echo(test_object)
+                    with self.assertRaises(pc_rpc.RemoteError):
+                        remote.non_existing_method()
+                    remote.quit()
+                finally:
+                    remote.close_rpc()
+            finally:
+                try:
+                    proc.wait(timeout=1)
+                except subprocess.TimeoutExpired:
+                    proc.kill()
+                    raise
+            self.assertEqual(test_object, test_object_back)
+
+
+class Echo(pc_rpc.WaitQuit):
+    def __init__(self):
+        pc_rpc.WaitQuit.__init__(self)
+
+    def echo(self, x):
+        return x
+
+
+def run_server():
+    loop = asyncio.get_event_loop()
+    try:
+        echo = Echo()
+        server = pc_rpc.Server(echo)
+        loop.run_until_complete(server.start(test_address, test_port))
+        try:
+            loop.run_until_complete(echo.wait_quit())
+        finally:
+            loop.run_until_complete(server.stop())
+    finally:
+        loop.close()
+
+if __name__ == "__main__":
+    run_server()