diff --git a/CHANGELOG.md b/CHANGELOG.md index 419671f..da32253 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/include/pyawaitable/awaitableobject.h b/include/pyawaitable/awaitableobject.h index 9567d90..e22ee83 100644 --- a/include/pyawaitable/awaitableobject.h +++ b/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 * diff --git a/src/_pyawaitable/awaitable.c b/src/_pyawaitable/awaitable.c index 25bc741..6fea0e3 100644 --- a/src/_pyawaitable/awaitable.c +++ b/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) { diff --git a/src/_pyawaitable/genwrapper.c b/src/_pyawaitable/genwrapper.c index 36fdeb5..9bf0614 100644 --- a/src/_pyawaitable/genwrapper.c +++ b/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; } diff --git a/src/_pyawaitable/mod.c b/src/_pyawaitable/mod.c index 10e9907..f5a1c45 100644 --- a/src/_pyawaitable/mod.c +++ b/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 diff --git a/src/pyawaitable/bindings.py b/src/pyawaitable/bindings.py index f6029ae..0169b90 100644 --- a/src/pyawaitable/bindings.py +++ b/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") diff --git a/src/pyawaitable/pyawaitable.h b/src/pyawaitable/pyawaitable.h index 6f6209b..5c62088 100644 --- a/src/pyawaitable/pyawaitable.h +++ b/src/pyawaitable/pyawaitable.h @@ -3,13 +3,14 @@ #include #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 diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py index 7229dd7..9f04508 100644 --- a/tests/test_callbacks.py +++ b/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