From 6fd17ff4d99d1dbcf98aa2b0ae96a5e23cdcdf16 Mon Sep 17 00:00:00 2001
From: Jason Paryani <github@jparyani.com>
Date: Wed, 4 Dec 2013 15:59:22 -0800
Subject: [PATCH] Add `then` function for VoidPromise

---
 capnp/capabilityHelper.h | 16 ++++++++++++++++
 capnp/capnp.pyx          | 22 ++++++++--------------
 capnp/capnp_cpp.pxd      |  1 +
 3 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/capnp/capabilityHelper.h b/capnp/capabilityHelper.h
index 23107ed..26708f0 100644
--- a/capnp/capabilityHelper.h
+++ b/capnp/capabilityHelper.h
@@ -70,6 +70,14 @@ PyObject * wrapPyFunc(PyObject * func, PyObject * arg) {
     return result;
 }
 
+PyObject * wrapPyFuncNoArg(PyObject * func) {
+    PyObject * result = PyObject_CallFunctionObjArgs(func, NULL);
+    Py_DECREF(func);
+
+    check_py_error();
+    return result;
+}
+
 void wrapRemoteCall(PyObject * func, capnp::Response<capnp::DynamicStruct> & arg) {
     wrap_remote_call(func, arg);
 
@@ -92,6 +100,14 @@ void wrapRemoteCall(PyObject * func, capnp::Response<capnp::DynamicStruct> & arg
                                      , [error_func](kj::Exception arg) { wrapPyFunc(error_func, wrap_kj_exception(arg)); } );
 }
 
+::kj::Promise<PyObject *> then(kj::Promise<void> & promise, PyObject * func, PyObject * error_func) {
+  if(error_func == Py_None)
+    return promise.then([func]() { return wrapPyFuncNoArg(func); } );
+  else
+    return promise.then([func]() { return wrapPyFuncNoArg(func); } 
+                                     , [error_func](kj::Exception arg) { return wrapPyFunc(error_func, wrap_kj_exception(arg)); } );
+}
+
 class PythonInterfaceDynamicImpl final: public capnp::DynamicCapability::Server {
 public:
   PyObject * py_server;
diff --git a/capnp/capnp.pyx b/capnp/capnp.pyx
index db3ab07..f8f5d69 100644
--- a/capnp/capnp.pyx
+++ b/capnp/capnp.pyx
@@ -1255,7 +1255,7 @@ 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')
 
-        ret = <object>self.thisptr.wait()
+        ret = <object>self.thisptr.wait() # TODO: make sure refcount is fine here...
         self.is_consumed = True
 
         return ret
@@ -1292,14 +1292,14 @@ cdef class _VoidPromise:
         self.is_consumed = True
 
 
-    # cpdef then(self, func, error_func=None) except +reraise_kj_exception:
-    #     if self.is_consumed:
-    #         raise RuntimeError('Promise was already used in a consuming operation. You can no longer use this Promise object')
+    cpdef then(self, func, error_func=None) except +reraise_kj_exception:
+        if self.is_consumed:
+            raise RuntimeError('Promise was already used in a consuming operation. You can no longer use this Promise object')
 
-    #     Py_INCREF(func)
-    #     Py_INCREF(error_func)
+        Py_INCREF(func)
+        Py_INCREF(error_func)
 
-    #     return Promise()._init(capnp.then(deref(self.thisptr), <PyObject *>func, <PyObject *>error_func))
+        return _Promise()._init(capnp.then(deref(self.thisptr), <PyObject *>func, <PyObject *>error_func))
 
 cdef class _RemotePromise:
     cdef RemotePromise * thisptr
@@ -1332,7 +1332,7 @@ cdef class _RemotePromise:
 
     cpdef then(self, func, error_func=None) except +reraise_kj_exception:
         if self.is_consumed:
-            raise ValueError('Promise was already used in a consuming operation. You can no longer use this Promise object')
+            raise RuntimeError('Promise was already used in a consuming operation. You can no longer use this Promise object')
 
         Py_INCREF(func)
         Py_INCREF(error_func)
@@ -1361,12 +1361,6 @@ cdef class _RemotePromise:
     def __dir__(self):
         return list(self.schema.fieldnames)
 
-    # def __str__(self):
-    #     return printStructReader(self.thisptr).flatten().cStr()
-
-    # def __repr__(self):
-    #     return '<%s reader %s>' % (self.schema.node.displayName, strStructReader(self.thisptr).cStr())
-
     def to_dict(self, verbose=False):
         return _to_dict(self, verbose)
 
diff --git a/capnp/capnp_cpp.pxd b/capnp/capnp_cpp.pxd
index 21e0764..aba305f 100644
--- a/capnp/capnp_cpp.pxd
+++ b/capnp/capnp_cpp.pxd
@@ -244,6 +244,7 @@ cdef extern from "capabilityHelper.h":
     # PyPromise there(EventLoop & loop, PyPromise & promise, PyObject * func, PyObject * error_func)
     PyPromise then(PyPromise & promise, PyObject * func, PyObject * error_func)
     VoidPromise then(RemotePromise & promise, PyObject * func, PyObject * error_func)
+    PyPromise then(VoidPromise & promise, PyObject * func, PyObject * error_func)
     cppclass PythonInterfaceDynamicImpl:
         PythonInterfaceDynamicImpl(PyObject *)
     DynamicCapability.Client new_client(InterfaceSchema&, PyObject *)