pyln: Add command notifications support

contrib/pyln-client/pyln/client/lightning.py

@@ -1,12 +1,13 @@
-from decimal import Decimal
-from math import floor, log10
-from typing import Optional, Union
 import json
 import logging
 import os
 import socket
 import warnings
+from contextlib import contextmanager
+from decimal import Decimal
 from json import JSONEncoder
+from math import floor, log10
+from typing import Optional, Union
 
 
 def _patched_default(self, obj):
@@ -283,8 +284,9 @@ def __init__(self, socket_path, executor=None, logger=logging, encoder_cls=json.
         self.decoder = decoder
         self.executor = executor
         self.logger = logger
+        self._notify = None
 
-        self.next_id = 0
+        self.next_id = 1
 
     def _writeobj(self, sock, obj):
         s = json.dumps(obj, ensure_ascii=False, cls=self.encoder_cls)
@@ -334,18 +336,44 @@ def call(self, method, payload=None):
         # FIXME: we open a new socket for every readobj call...
         sock = UnixSocket(self.socket_path)
         this_id = self.next_id
-        self._writeobj(sock, {
+        self.next_id += 0
+        buf = b''
+
+        if self._notify is not None:
+            # Opt into the notifications support
+            self._writeobj(sock, {
+                "jsonrpc": "2.0",
+                "method": "notifications",
+                "id": 0,
+                "params": {
+                    "enable": True
+                },
+            })
+            _, buf = self._readobj(sock, buf)
+
+        request = {
             "jsonrpc": "2.0",
             "method": method,
             "params": payload,
             "id": this_id,
-        })
-        self.next_id += 1
-        buf = b''
+        }
+
+        self._writeobj(sock, request)
         while True:
             resp, buf = self._readobj(sock, buf)
-            # FIXME: We should offer a callback for notifications.
-            if 'method' not in resp or 'id' in resp:
+            id = resp.get("id", None)
+            meth = resp.get("method", None)
+
+            if meth == 'message' and self._notify is not None:
+                n = resp['params']
+                self._notify(
+                    message=n.get('message', None),
+                    progress=n.get('progress', None),
+                    request=request
+                )
+                continue
+
+            if meth is None or id is None:
                 break
 
         self.logger.debug("Received response for %s call: %r", method, resp)
@@ -361,6 +389,31 @@ def call(self, method, payload=None):
             raise ValueError("Malformed response, \"result\" missing.")
         return resp["result"]
 
+    @contextmanager
+    def notify(self, fn):
+        """Register a notification callback to use for a set of RPC calls.
+
+        This is a context manager and should be used like this:
+
+        ```python
+        def fn(message, progress, request, **kwargs):
+            print(message)
+
+        with rpc.notify(fn):
+            rpc.somemethod()
+        ```
+
+        The `fn` function will be called once for each notification
+        the is sent by `somemethod`. This is a context manager,
+        meaning that multiple commands can share the same context, and
+        the same notification function.
+
+        """
+        old = self._notify
+        self._notify = fn
+        yield
+        self._notify = old
+
 
 class LightningRpc(UnixDomainSocketRpc):
     """

contrib/pyln-client/pyln/client/plugin.py

@@ -141,6 +141,47 @@ def set_exception(self, exc: Union[Exception, RpcException]) -> None:
     def _write_result(self, result: dict) -> None:
         self.plugin._write_locked(result)
 
+    def _notify(self, method: str, params: JSONType) -> None:
+        """Send a notification to the caller.
+
+        Can contain a variety of things, but is usually used to report
+        progress or command status.
+
+        """
+        self._write_result({
+            'jsonrpc': '2.0',
+            'params': params,
+            "method": method,
+        })
+
+    def notify(self, message: str, level: str = 'info') -> None:
+        """Send a message notification to the caller.
+        """
+        self._notify(
+            "message",
+            params={
+                'id': self.id,
+                'level': level,
+                'message': message,
+            }
+        )
+
+    def progress(self,
+                 progress: int,
+                 total: int,
+                 stage: Optional[int] = None,
+                 stage_total: Optional[int] = None
+                 ) -> None:
+        d: Dict[str, JSONType] = {
+            "id": self.id,
+            "num": progress,
+            "total": total,
+        }
+        if stage is not None and stage_total is not None:
+            d['stage'] = {"num": stage, "total": stage_total}
+
+        self._notify("progress", d)
+
 
 # If a hook call fails we need to coerce it into something the main daemon can
 # handle. Returning an error is not an option since we explicitly do not allow
@@ -639,22 +680,14 @@ def log(self, message: str, level: str = 'info') -> None:
     def notify_message(self, request: Request, message: str,
                        level: str = 'info') -> None:
         """Send a notification message to sender of this request"""
-        self.notify("message", {"id": request.id,
-                                "level": level,
-                                "message": message})
+        request.notify(message=message)
 
     def notify_progress(self, request: Request,
                         progress: int, progress_total: int,
                         stage: Optional[int] = None,
                         stage_total: Optional[int] = None) -> None:
-        """Send a progerss message to sender of this request: if more than one stage, set stage and stage_total"""
-        d: Dict[str, Any] = {"id": request.id,
-                             "num": progress,
-                             "total": progress_total}
-        if stage_total is not None:
-            d['stage'] = {"num": stage, "total": stage_total}
-
-        self.notify("progress", d)
+        """Send a progress message to sender of this request: if more than one stage, set stage and stage_total"""
+        request.progress(progress, progress_total, stage, stage_total)
 
     def _parse_request(self, jsrequest: Dict[str, JSONType]) -> Request:
         i = jsrequest.get('id', None)

tests/plugins/countdown.py

@@ -0,0 +1,19 @@
+#!/usr/bin/env python3
+
+from pyln.client import Plugin
+import time
+
+plugin = Plugin()
+
+
+@plugin.method("countdown")
+def countdown(count, plugin, request):
+    count = int(count)
+    for i in range(count):
+        time.sleep(0.1)
+        request.notify("{}/{}".format(i, count), "INFO")
+
+    return "Done"
+
+
+plugin.run()

tests/test_plugin.py

@@ -2027,15 +2027,18 @@ def test_notify(node_factory):
             assert out[2 + i].endswith("|\n")
         else:
             assert out[2 + i].endswith("|\r")
-    assert out[102] == '\r'
+
+    assert out[102] == '# Beginning stage 2\n'
+    assert out[103] == '\r'
+
     for i in range(10):
-        assert out[103 + i].startswith("# Stage 2/2 {:>2}/10 |".format(1 + i))
+        assert out[104 + i].startswith("# Stage 2/2 {:>2}/10 |".format(1 + i))
         if i == 9:
-            assert out[103 + i].endswith("|\n")
+            assert out[104 + i].endswith("|\n")
         else:
-            assert out[103 + i].endswith("|\r")
-    assert out[113] == '"This worked"\n'
-    assert len(out) == 114
+            assert out[104 + i].endswith("|\r")
+    assert out[114] == '"This worked"\n'
+    assert len(out) == 115
 
     # At debug level, we get the second prompt.
     out = subprocess.check_output(['cli/lightning-cli',
@@ -2199,3 +2202,28 @@ def test_dynamic_args(node_factory):
     l1.rpc.plugin_stop(plugin_path)
 
     assert [p for p in l1.rpc.listconfigs()['plugins'] if p['path'] == plugin_path] == []
+
+
+def test_pyln_request_notify(node_factory):
+    """Test that pyln-client plugins can send notifications.
+    """
+    plugin_path = os.path.join(
+        os.path.dirname(__file__), 'plugins/countdown.py'
+    )
+    l1 = node_factory.get_node(options={'plugin': plugin_path})
+    notifications = []
+
+    def n(*args, message, **kwargs):
+        print("Got a notification:", message)
+        notifications.append(message)
+
+    with l1.rpc.notify(n):
+        l1.rpc.countdown(10)
+
+    expected = ['{}/10'.format(i) for i in range(10)]
+    assert expected == notifications
+
+    # Calling without the context manager we should not get any notifications
+    notifications = []
+    l1.rpc.countdown(10)
+    assert notifications == []