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/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