From 979d24e8c42925b54c2bb1731935dcdc2b633086 Mon Sep 17 00:00:00 2001
From: Leonid Shamis <leonid.shamis@gmail.com>
Date: Thu, 13 Jan 2022 09:59:19 -0800
Subject: [PATCH] Adding decorators

---
 a0.py             | 97 +++++++++++++++++++++++++++++++++++++++++++++++
 tests/test_dec.py | 38 +++++++++++++++++++
 tests/test_rpc.py |  2 +
 3 files changed, 137 insertions(+)
 create mode 100644 tests/test_dec.py

diff --git a/a0.py b/a0.py
index 83d4e29..d096624 100644
--- a/a0.py
+++ b/a0.py
@@ -1,6 +1,7 @@
 from alephzero_bindings import *
 import asyncio
 import base64
+import concurrent.futures
 import json
 import jsonpointer
 import threading
@@ -110,6 +111,19 @@ def update_configs():
         cfg.registry = next_reg
 
 
+################
+# Nice-to-Have #
+################
+
+
+def _rcp_client_send_fut(client, *args):
+    fut = concurrent.futures.Future()
+    client.send(*args, lambda pkt: fut.set_result(pkt))
+    return fut
+
+
+RpcClient.send_fut = _rcp_client_send_fut
+
 ###########
 # AsyncIO #
 ###########
@@ -235,6 +249,89 @@ def onloop():
         return await ns.fut
 
 
+##############
+# Decorators #
+##############
+
+
+class dec:
+
+    class Entry:
+
+        def __init__(self, fn, topic, factory):
+            self.fn = fn
+            self.topic = topic
+            self.factory = factory
+            self._obj = None
+
+        def running(self):
+            return self._obj is not None
+
+        def start(self):
+            if self.running():
+                return
+            self._obj = self.factory()
+
+        def stop(self):
+            self._obj = None
+
+    registry = []
+
+    @staticmethod
+    def start_all():
+        for entry in dec.registry:
+            entry.start()
+
+    @staticmethod
+    def stop_all():
+        for entry in dec.registry:
+            entry.stop()
+
+    # Example:
+    #
+    # @a0.dec.sub("foo")
+    # def foo(pkt):
+    #     ...
+    @staticmethod
+    def sub(*args, topic=None, opts=None, init_=None, iter_=None):
+        if topic is None and args and type(args[0]) is str:
+            topic = args[0]
+
+        def decorator(fn):
+            topic_ = topic or fn.__name__
+            opts_ = _aio_read_base._make_opts(opts, init_, iter_)
+            dec.registry.append(
+                dec.Entry(fn, topic_, lambda: Subscriber(topic_, opts_, fn)))
+            return fn
+
+        if len(args) == 1 and callable(args[0]):
+            return decorator(args[0])
+
+        return decorator
+
+    # Example:
+    #
+    # @a0.dec.rpc("foo")
+    # def foo(req):
+    #     ...
+    #     return resp
+    @staticmethod
+    def rpc(*args, topic=None):
+        if topic is None and args and type(args[0]) is str:
+            topic = args[0]
+
+        def decorator(fn):
+            topic_ = topic or fn.__name__
+            dec.registry.append(
+                dec.Entry(fn, topic_, lambda: RpcServer(topic_, fn, None)))
+            return fn
+
+        if len(args) == 1 and callable(args[0]):
+            return decorator(args[0])
+
+        return decorator
+
+
 ##########
 # Remote #
 ##########
diff --git a/tests/test_dec.py b/tests/test_dec.py
new file mode 100644
index 0000000..89d60df
--- /dev/null
+++ b/tests/test_dec.py
@@ -0,0 +1,38 @@
+import a0
+import threading
+
+
+def test_dec():
+
+    cv = threading.Condition()
+    payloads = []
+
+    @a0.dec.sub("foo")
+    def foo(pkt):
+        with cv:
+            payloads.append(pkt.payload)
+            cv.notify()
+
+    @a0.dec.sub
+    def bar(pkt):
+        with cv:
+            payloads.append(pkt.payload)
+            cv.notify()
+
+    @a0.dec.sub("bar")
+    def bar2(pkt):
+        with cv:
+            payloads.append(pkt.payload)
+            cv.notify()
+
+    a0.dec.start_all()
+
+    foo_pub = a0.Publisher("foo")
+    foo_pub.pub("hello")
+
+    bar_pub = a0.Publisher("bar")
+    bar_pub.pub("world")
+
+    with cv:
+        cv.wait_for(lambda: len(payloads) == 3)
+    assert sorted(payloads) == [b"hello", b"world", b"world"]
diff --git a/tests/test_rpc.py b/tests/test_rpc.py
index db25d96..7130b08 100644
--- a/tests/test_rpc.py
+++ b/tests/test_rpc.py
@@ -56,3 +56,5 @@ def onreply(pkt):
     assert reply.payload == b"slept"
     with pytest.raises(RuntimeError, match="Connection timed out"):
         client.send_blocking("sleep", timeout=a0.TimeMono.now() + 0.1)
+
+    assert client.send_fut("reply").result().payload == b"reply"