[TASK]-[PYTHON-5623]: Change with_transaction callback return type to Awaitable (#2594)

Co-authored-by: Logan Pulley <logan@pulley.host>
(cherry picked from commit 5f00966)

Author: 2 people

doc/changelog.rst

@@ -1,6 +1,22 @@
 Changelog
 =========
 
+Changes in Version 4.15.4 (2025/11/11)
+--------------------------------------
+
+Version 4.15.4 is a bug fix release.
+
+- Relaxed the callback type of :meth:`~pymongo.asynchronous.client_session.AsyncClientSession.with_transaction` to allow the broader Awaitable type rather than only Coroutine objects.
+- Added the missing Python 3.14 trove classifier to the package metadata.
+
+Issues Resolved
+...............
+
+See the `PyMongo 4.15.4 release notes in JIRA`_ for the list of resolved issues
+in this release.
+
+.. _PyMongo 4.15.4 release notes in JIRA: https://jira.mongodb.org/secure/ReleaseNote.jspa?projectId=10004&version=47237
+
 Changes in Version 4.15.3 (2025/10/07)
 --------------------------------------
 

pymongo/asynchronous/client_session.py

@@ -143,8 +143,8 @@
     TYPE_CHECKING,
     Any,
     AsyncContextManager,
+    Awaitable,
     Callable,
-    Coroutine,
     Mapping,
     MutableMapping,
     NoReturn,
@@ -600,7 +600,7 @@ def _inherit_option(self, name: str, val: _T) -> _T:
 
     async def with_transaction(
         self,
-        callback: Callable[[AsyncClientSession], Coroutine[Any, Any, _T]],
+        callback: Callable[[AsyncClientSession], Awaitable[_T]],
         read_concern: Optional[ReadConcern] = None,
         write_concern: Optional[WriteConcern] = None,
         read_preference: Optional[_ServerMode] = None,

test/asynchronous/test_transactions.py

@@ -15,6 +15,7 @@
 """Execute Transactions Spec tests."""
 from __future__ import annotations
 
+import asyncio
 import sys
 from io import BytesIO
 from test.asynchronous.utils_spec_runner import AsyncSpecRunner
@@ -469,6 +470,17 @@ async def callback2(session):
         async with self.client.start_session() as s:
             self.assertEqual(await s.with_transaction(callback2), "Foo")
 
+    @async_client_context.require_transactions
+    @async_client_context.require_async
+    async def test_callback_awaitable_no_coroutine(self):
+        def callback(_):
+            future = asyncio.Future()
+            future.set_result("Foo")
+            return future
+
+        async with self.client.start_session() as s:
+            self.assertEqual(await s.with_transaction(callback), "Foo")
+
     @async_client_context.require_transactions
     async def test_callback_not_retried_after_timeout(self):
         listener = OvertCommandListener()

test/test_transactions.py

@@ -15,6 +15,7 @@
 """Execute Transactions Spec tests."""
 from __future__ import annotations
 
+import asyncio
 import sys
 from io import BytesIO
 from test.utils_spec_runner import SpecRunner
@@ -461,6 +462,17 @@ def callback2(session):
         with self.client.start_session() as s:
             self.assertEqual(s.with_transaction(callback2), "Foo")
 
+    @client_context.require_transactions
+    @client_context.require_async
+    def test_callback_awaitable_no_coroutine(self):
+        def callback(_):
+            future = asyncio.Future()
+            future.set_result("Foo")
+            return future
+
+        with self.client.start_session() as s:
+            self.assertEqual(s.with_transaction(callback), "Foo")
+
     @client_context.require_transactions
     def test_callback_not_retried_after_timeout(self):
         listener = OvertCommandListener()

tools/synchro.py

@@ -322,6 +322,14 @@ def translate_coroutine_types(lines: list[str]) -> list[str]:
             index = lines.index(type)
             new = type.replace(old, res.group(3))
             lines[index] = new
+    coroutine_types = [line for line in lines if "Awaitable[" in line]
+    for type in coroutine_types:
+        res = re.search(r"Awaitable\[([A-z]+)\]", type)
+        if res:
+            old = res[0]
+            index = lines.index(type)
+            new = type.replace(old, res.group(1))
+            lines[index] = new
     return lines