From cfb60cddf0ddf27ee3e06ae2616e23366776b712 Mon Sep 17 00:00:00 2001
From: Jason Paryani <github@jparyani.com>
Date: Sun, 13 Apr 2014 18:26:01 -0700
Subject: [PATCH] Add arguments checking to `then` callback

---
 capnp/lib/capnp.pyx     | 11 +++++++++++
 test/test_capability.py | 10 ++++++++++
 2 files changed, 21 insertions(+)

diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx
index b6ad05a..71262c6 100644
--- a/capnp/lib/capnp.pyx
+++ b/capnp/lib/capnp.pyx
@@ -1463,6 +1463,17 @@ cdef class Promise:
         if self.is_consumed:
             raise ValueError('Promise was already used in a consuming operation. You can no longer use this Promise object')
 
+        argspec = None
+        try:
+            argspec = _inspect.getargspec(func)
+        except:
+            pass
+        if argspec:
+            args_length = len(argspec.args) if argspec.args else 0
+            defaults_length = len(argspec.defaults) if argspec.defaults else 0
+            if args_length - defaults_length != 1:
+                raise ValueError('Function passed to `then` call must take exactly one argument')
+
         self.is_consumed = True
 
         return Promise()._init(helpers.then(deref(self.thisptr), <PyObject *>func, <PyObject *>error_func).attach(capnp.makePyRefCounter(<PyObject *>func), capnp.makePyRefCounter(<PyObject *>error_func)), self)
diff --git a/test/test_capability.py b/test/test_capability.py
index 6f1e79d..7278f8d 100644
--- a/test/test_capability.py
+++ b/test/test_capability.py
@@ -289,3 +289,13 @@ def test_timer():
     joined = capnp.join_promises([promise, canceller])
     with pytest.raises(Exception):
         joined.wait()
+
+
+def test_then_args():
+    capnp.Promise(0).then(lambda x: 1)
+
+    with pytest.raises(ValueError):
+        capnp.Promise(0).then(lambda: 1)
+
+    with pytest.raises(ValueError):
+        capnp.Promise(0).then(lambda x, y: 1)