Add PyAwaitable_DeferAwait

CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Objects returned by a PyAwaitable object's `__await__` are now garbage collected (*i.e.*, they don't leak with rare circular references).
 - Removed limit on number of stored callbacks or values.
 - Switched some user-error messages to `RuntimeError` instead of `SystemError`.
+- Added `PyAwaitable_DeferAwait` for executing code when the awaitable object is called by the event loop.
 
 ## [1.3.0] - 2024-10-26
 

include/pyawaitable/awaitableobject.h

@@ -8,6 +8,7 @@
 
 typedef int (*awaitcallback)(PyObject *, PyObject *);
 typedef int (*awaitcallback_err)(PyObject *, PyObject *);
+typedef int (*defer_callback)(PyObject *);
 
 typedef struct _pyawaitable_callback
 {
@@ -52,6 +53,8 @@ int pyawaitable_await_impl(
     awaitcallback_err err
 );
 
+int pyawaitable_defer_await_impl(PyObject *aw, defer_callback cb);
+
 void pyawaitable_cancel_impl(PyObject *aw);
 
 PyObject *

src/_pyawaitable/awaitable.c

@@ -172,6 +172,32 @@ pyawaitable_await_impl(
     return 0;
 }
 
+int
+pyawaitable_defer_await_impl(PyObject *awaitable, defer_callback cb)
+{
+    PyAwaitableObject *aw = (PyAwaitableObject *) awaitable;
+    pyawaitable_callback *aw_c = PyMem_Malloc(sizeof(pyawaitable_callback));
+    if (aw_c == NULL)
+    {
+        PyErr_NoMemory();
+        return -1;
+    }
+
+    aw_c->coro = NULL;
+    aw_c->callback = (awaitcallback)cb;
+    aw_c->err_callback = NULL;
+    aw_c->done = false;
+
+    if (pyawaitable_array_append(&aw->aw_callbacks, aw_c) < 0)
+    {
+        PyMem_Free(aw_c);
+        PyErr_NoMemory();
+        return -1;
+    }
+
+    return 0;
+}
+
 int
 pyawaitable_set_result_impl(PyObject *awaitable, PyObject *result)
 {

src/_pyawaitable/genwrapper.c

@@ -167,6 +167,31 @@ genwrapper_next(PyObject *self)
             );
         }
 
+        if (cb->callback != NULL && cb->coro == NULL)
+        {
+            int def_res = ((defer_callback)cb->callback)((PyObject*)aw);
+
+            // If we recently cancelled, then cb is no longer valid
+            if (aw->aw_recently_cancelled)
+            {
+              cb = NULL;
+            }
+
+            if (def_res < 0 && !PyErr_Occurred())
+            {
+                PyErr_SetString(
+                    PyExc_SystemError,
+                    "pyawaitable: callback returned -1 without exception set"
+                );
+                DONE_IF_OK(cb);
+                return NULL;
+            }
+
+            // Callback is done.
+            DONE_IF_OK(cb);
+            return genwrapper_next(self);
+        }
+
         if (
             Py_TYPE(cb->coro)->tp_as_async == NULL ||
             Py_TYPE(cb->coro)->tp_as_async->am_await == NULL
@@ -311,7 +336,7 @@ genwrapper_next(PyObject *self)
                 PyExc_RuntimeError,
                 "pyawaitable: user callback returned -1 without exception set"
             );
-            DONE(cb);
+            DONE_IF_OK(cb);
             AW_DONE();
             return NULL;
         }

src/_pyawaitable/mod.c

@@ -55,7 +55,8 @@ static PyAwaitableABI _abi_interface =
     pyawaitable_get_impl,
     pyawaitable_get_arb_impl,
     pyawaitable_get_int_impl,
-    pyawaitable_async_with_impl
+    pyawaitable_async_with_impl,
+    pyawaitable_defer_await_impl
 };
 
 PyMODINIT_FUNC

src/pyawaitable/bindings.py

@@ -6,7 +6,7 @@
 
 from . import abi
 
-__all__ = "abi", "add_await", "awaitcallback", "awaitcallback_err"
+__all__ = ["abi", "add_await", "awaitcallback", "awaitcallback_err", "defer_callback", "defer_await"]
 
 get_pointer = pythonapi.PyCapsule_GetPointer
 get_pointer.argtypes = (ctypes.py_object, ctypes.c_void_p)
@@ -48,6 +48,8 @@ def __getattribute__(self, name: str) -> Any:
     ctypes.c_int, ctypes.py_object, ctypes.py_object
 )
 awaitcallback_err = awaitcallback
+defer_callback = ctypes.PYFUNCTYPE(
+    ctypes.c_int, ctypes.py_object)
 
 
 class AwaitableABI(PyABI):
@@ -163,8 +165,17 @@ class AwaitableABI(PyABI):
                 awaitcallback_err,
             ),
         ),
+        (
+            "defer_await",
+            ctypes.PYFUNCTYPE(
+                ctypes.c_int,
+                ctypes.py_object,
+                defer_callback
+            ),
+        ),
     ]
 
 
 abi = AwaitableABI.from_capsule(abi.v1)
 add_await = getattr(abi, "await")
+defer_await = getattr(abi, "defer_await")

src/pyawaitable/pyawaitable.h

@@ -3,13 +3,14 @@
 #include <Python.h>
 
 #define PYAWAITABLE_MAJOR_VERSION 1
-#define PYAWAITABLE_MINOR_VERSION 3
+#define PYAWAITABLE_MINOR_VERSION 4
 #define PYAWAITABLE_MICRO_VERSION 0
 /* Per CPython Conventions: 0xA for alpha, 0xB for beta, 0xC for release candidate or 0xF for final. */
 #define PYAWAITABLE_RELEASE_LEVEL 0xF
 
 typedef int (*awaitcallback)(PyObject *, PyObject *);
 typedef int (*awaitcallback_err)(PyObject *, PyObject *);
+typedef int (*defer_callback)(PyObject *);
 
 typedef struct _PyAwaitableObject PyAwaitableObject;
 
@@ -52,6 +53,9 @@ typedef struct _pyawaitable_abi
         awaitcallback cb,
         awaitcallback_err err
     );
+    int (*defer_await)(
+        PyObject *aw,
+        defer_callback cb);
 } PyAwaitableABI;
 
 #ifdef PYAWAITABLE_THIS_FILE_INIT
@@ -67,6 +71,7 @@ extern PyAwaitableABI *pyawaitable_abi;
 #define pyawaitable_await pyawaitable_abi->await
 #define pyawaitable_await_function pyawaitable_abi->await_function
 #define pyawaitable_async_with pyawaitable_abi->async_with
+#define pyawaitable_defer_await pyawaitable_abi->defer_await
 
 #define pyawaitable_save pyawaitable_abi->save
 #define pyawaitable_save_arb pyawaitable_abi->save_arb
@@ -127,6 +132,7 @@ pyawaitable_init()
 #define PyAwaitable_AddAwait pyawaitable_await
 #define PyAwaitable_AwaitFunction pyawaitable_await_function
 #define PyAwaitable_AsyncWith pyawaitable_async_with
+#define PyAwaitable_DeferAwait pyawaitable_defer_await
 
 #define PyAwaitable_SaveValues pyawaitable_save
 #define PyAwaitable_SaveArbValues pyawaitable_save_arb

tests/test_callbacks.py

@@ -5,7 +5,8 @@
 
 import pyawaitable
 from pyawaitable.bindings import (abi, add_await, awaitcallback,
-                                  awaitcallback_err)
+                                  awaitcallback_err, defer_callback,
+                                  defer_await)
 
 
 @limit_leaks
@@ -172,3 +173,32 @@ def callback(awaitable: pyawaitable.PyAwaitable, result: None) -> int:
     await awaitable
     assert called == amount
     assert awaited == amount
+
+@limit_leaks
+@pytest.mark.asyncio
+async def test_deferred_await():
+    called = 0
+    awaitable = abi.new()
+
+    async def coro(value: int):
+        await asyncio.sleep(0)
+        return value * 2
+
+    @awaitcallback
+    def cb(awaitable_inner: pyawaitable.PyAwaitable, result: int) -> int:
+        assert awaitable_inner is awaitable
+        assert result == 42
+        nonlocal called
+        called += 1
+        return 0
+
+    @defer_callback
+    def defer_cb(awaitable: pyawaitable.PyAwaitable):
+        nonlocal called
+        called += 1
+        add_await(awaitable, coro(21), cb, awaitcallback_err(0))
+        return 0
+
+    defer_await(awaitable, defer_cb)
+    await awaitable
+    assert called == 2