Async cursors' for_each accepts coroutines as well (#339)

Author: hemidactylus

astrapy/data/cursor.py

@@ -19,7 +19,8 @@
 from copy import deepcopy
 from decimal import Decimal
 from enum import Enum
-from typing import Any, Callable, Generic, TypeVar, cast
+from inspect import iscoroutinefunction
+from typing import Any, Awaitable, Callable, Generic, TypeVar, cast
 
 from typing_extensions import override
 
@@ -1106,7 +1107,7 @@ def for_each(
         results in the cursor entering CLOSED state.
 
         Args:
-            function: a callback function whose only parameter is the type returned
+            function: a callback function whose only parameter is of the type returned
                 by the cursor. This callback is invoked once per each document yielded
                 by the cursor. If the callback returns a `False`, the `for_each`
                 invocation stops early and returns without consuming further documents.
@@ -1729,14 +1730,14 @@ def _composite(document: TRAW) -> TNEW:
 
     async def for_each(
         self,
-        function: Callable[[T], bool | None],
+        function: Callable[[T], bool | None] | Callable[[T], Awaitable[bool | None]],
         *,
         general_method_timeout_ms: int | None = None,
         timeout_ms: int | None = None,
     ) -> None:
         """
         Consume the remaining documents in the cursor, invoking a provided callback
-        function on each of them.
+        function -- or coroutine -- on each of them.
 
         Calling this method on a CLOSED cursor results in an error.
 
@@ -1751,8 +1752,9 @@ async def for_each(
         adaptations to the async interface.
 
         Args:
-            function: a callback function whose only parameter is the type returned
-                by the cursor. This callback is invoked once per each document yielded
+            function: a callback function, or a coroutine, whose only parameter is of
+                the type returned by the cursor.
+                This callback is invoked once per each document yielded
                 by the cursor. If the callback returns a `False`, the `for_each`
                 invocation stops early and returns without consuming further documents.
             general_method_timeout_ms: a timeout, in milliseconds, for the whole
@@ -1772,8 +1774,12 @@ async def for_each(
             overall_timeout_ms=copy_ovr_ms,
         )
         self._imprint_internal_state(_cursor)
+        is_coro = iscoroutinefunction(function)
         async for document in _cursor:
-            res = function(document)
+            if is_coro:
+                res = await function(document)  # type: ignore[misc]
+            else:
+                res = function(document)
             if res is False:
                 break
         _cursor._imprint_internal_state(self)
@@ -2384,7 +2390,7 @@ def for_each(
         results in the cursor entering CLOSED state.
 
         Args:
-            function: a callback function whose only parameter is the type returned
+            function: a callback function whose only parameter is of the type returned
                 by the cursor. This callback is invoked once per each row yielded
                 by the cursor. If the callback returns a `False`, the `for_each`
                 invocation stops early and returns without consuming further rows.
@@ -2445,8 +2451,8 @@ def for_each(
             overall_timeout_ms=copy_ovr_ms,
         )
         self._imprint_internal_state(_cursor)
-        for document in _cursor:
-            res = function(document)
+        for row in _cursor:
+            res = function(row)
             if res is False:
                 break
         _cursor._imprint_internal_state(self)
@@ -3006,14 +3012,14 @@ def _composite(document: TRAW) -> TNEW:
 
     async def for_each(
         self,
-        function: Callable[[T], bool | None],
+        function: Callable[[T], bool | None] | Callable[[T], Awaitable[bool | None]],
         *,
         general_method_timeout_ms: int | None = None,
         timeout_ms: int | None = None,
     ) -> None:
         """
         Consume the remaining rows in the cursor, invoking a provided callback
-        function on each of them.
+        function -- or coroutine -- on each of them.
 
         Calling this method on a CLOSED cursor results in an error.
 
@@ -3028,8 +3034,9 @@ async def for_each(
         adaptations to the async interface.
 
         Args:
-            function: a callback function whose only parameter is the type returned
-                by the cursor. This callback is invoked once per each row yielded
+            function: a callback function, or a coroutine, whose only parameter is of
+                the type returned by the cursor.
+                This callback is invoked once per each row yielded
                 by the cursor. If the callback returns a `False`, the `for_each`
                 invocation stops early and returns without consuming further rows.
             general_method_timeout_ms: a timeout, in milliseconds, for the whole
@@ -3049,8 +3056,12 @@ async def for_each(
             overall_timeout_ms=copy_ovr_ms,
         )
         self._imprint_internal_state(_cursor)
-        async for document in _cursor:
-            res = function(document)
+        is_coro = iscoroutinefunction(function)
+        async for row in _cursor:
+            if is_coro:
+                res = await function(row)  # type: ignore[misc]
+            else:
+                res = function(row)
             if res is False:
                 break
         _cursor._imprint_internal_state(self)

tests/base/integration/collections/test_collection_cursor_async.py

@@ -314,6 +314,20 @@ def marker0(row: dict[str, Any], acc: list[dict[str, Any]] = accum0) -> None:
         assert accum0 == base_rows
         assert fe_cur.state == FindCursorState.CLOSED
 
+        # full for_each, coroutine
+        aaccum0: list[dict[str, Any]] = []
+
+        async def amarker0(
+            row: dict[str, Any],
+            acc: list[dict[str, Any]] = aaccum0,
+        ) -> None:
+            acc += [row]
+
+        afe_cur = async_filled_collection.find()
+        await afe_cur.for_each(amarker0)
+        assert aaccum0 == base_rows
+        assert afe_cur.state == FindCursorState.CLOSED
+
         # partially-consumed for_each
         accum1: list[dict[str, Any]] = []
 
@@ -327,6 +341,22 @@ def marker1(row: dict[str, Any], acc: list[dict[str, Any]] = accum1) -> None:
         assert accum1 == base_rows[11:]
         assert pfe_cur.state == FindCursorState.CLOSED
 
+        # partially-consumed for_each, coroutine
+        aaccum1: list[dict[str, Any]] = []
+
+        async def amarker1(
+            row: dict[str, Any],
+            acc: list[dict[str, Any]] = aaccum1,
+        ) -> None:
+            acc += [row]
+
+        apfe_cur = async_filled_collection.find()
+        for _ in range(11):
+            await apfe_cur.__anext__()
+        await apfe_cur.for_each(amarker1)
+        assert aaccum1 == base_rows[11:]
+        assert apfe_cur.state == FindCursorState.CLOSED
+
         # mapped for_each
         accum2: list[int] = []
 
@@ -340,6 +370,19 @@ def marker2(val: int, acc: list[int] = accum2) -> None:
         assert accum2 == [mint(row) for row in base_rows[17:]]
         assert mfe_cur.state == FindCursorState.CLOSED
 
+        # mapped for_each, coroutine
+        aaccum2: list[int] = []
+
+        async def amarker2(val: int, acc: list[int] = aaccum2) -> None:
+            acc += [val]
+
+        amfe_cur = async_filled_collection.find().map(mint)
+        for _ in range(17):
+            await amfe_cur.__anext__()
+        await amfe_cur.for_each(amarker2)
+        assert aaccum2 == [mint(row) for row in base_rows[17:]]
+        assert amfe_cur.state == FindCursorState.CLOSED
+
         # breaking (early) for_each
         accum3: list[dict[str, Any]] = []
 
@@ -354,6 +397,23 @@ def marker3(row: dict[str, Any], acc: list[dict[str, Any]] = accum3) -> bool:
         bfe_another = await bfe_cur.__anext__()
         assert bfe_another == base_rows[5]
 
+        # breaking (early) for_each, coroutine
+        aaccum3: list[dict[str, Any]] = []
+
+        async def amarker3(
+            row: dict[str, Any],
+            acc: list[dict[str, Any]] = aaccum3,
+        ) -> bool:
+            acc += [row]
+            return len(acc) < 5
+
+        abfe_cur = async_filled_collection.find()
+        await abfe_cur.for_each(amarker3)
+        assert aaccum3 == base_rows[:5]
+        assert abfe_cur.state == FindCursorState.STARTED
+        abfe_another = await abfe_cur.__anext__()
+        assert abfe_another == base_rows[5]
+
         # nonbool-nonbreaking for_each
         accum4: list[dict[str, Any]] = []
 
@@ -366,6 +426,21 @@ def marker4(row: dict[str, Any], acc: list[dict[str, Any]] = accum4) -> int:
         assert accum4 == base_rows
         assert nbfe_cur.state == FindCursorState.CLOSED
 
+        # nonbool-nonbreaking for_each, coroutine
+        aaccum4: list[dict[str, Any]] = []
+
+        async def amarker4(
+            row: dict[str, Any],
+            acc: list[dict[str, Any]] = aaccum4,
+        ) -> int:
+            acc += [row]
+            return 8 if len(acc) < 5 else 0
+
+        anbfe_cur = async_filled_collection.find()
+        await anbfe_cur.for_each(amarker4)  # type: ignore[arg-type]
+        assert aaccum4 == base_rows
+        assert anbfe_cur.state == FindCursorState.CLOSED
+
     @pytest.mark.describe("test of collection cursors, serdes options obeyance, async")
     async def test_collection_cursors_serdes_options_async(
         self,

tests/base/integration/tables/test_table_cursor_async.py

@@ -314,6 +314,20 @@ def marker0(row: dict[str, Any], acc: list[dict[str, Any]] = accum0) -> None:
         assert accum0 == base_rows
         assert fe_cur.state == FindCursorState.CLOSED
 
+        # full for_each, coroutine
+        aaccum0: list[dict[str, Any]] = []
+
+        async def amarker0(
+            row: dict[str, Any],
+            acc: list[dict[str, Any]] = aaccum0,
+        ) -> None:
+            acc += [row]
+
+        afe_cur = filled_composite_atable.find()
+        await afe_cur.for_each(amarker0)
+        assert aaccum0 == base_rows
+        assert afe_cur.state == FindCursorState.CLOSED
+
         # partially-consumed for_each
         accum1: list[dict[str, Any]] = []
 
@@ -327,6 +341,22 @@ def marker1(row: dict[str, Any], acc: list[dict[str, Any]] = accum1) -> None:
         assert accum1 == base_rows[11:]
         assert pfe_cur.state == FindCursorState.CLOSED
 
+        # partially-consumed for_each, coroutine
+        aaccum1: list[dict[str, Any]] = []
+
+        async def amarker1(
+            row: dict[str, Any],
+            acc: list[dict[str, Any]] = aaccum1,
+        ) -> None:
+            acc += [row]
+
+        apfe_cur = filled_composite_atable.find()
+        for _ in range(11):
+            await apfe_cur.__anext__()
+        await apfe_cur.for_each(amarker1)
+        assert aaccum1 == base_rows[11:]
+        assert apfe_cur.state == FindCursorState.CLOSED
+
         # mapped for_each
         accum2: list[int] = []
 
@@ -340,6 +370,22 @@ def marker2(val: int, acc: list[int] = accum2) -> None:
         assert accum2 == [mint(row) for row in base_rows[17:]]
         assert mfe_cur.state == FindCursorState.CLOSED
 
+        # mapped for_each, coroutine
+        aaccum2: list[int] = []
+
+        async def amarker2(
+            val: int,
+            acc: list[int] = aaccum2,
+        ) -> None:
+            acc += [val]
+
+        amfe_cur = filled_composite_atable.find().map(mint)
+        for _ in range(17):
+            await amfe_cur.__anext__()
+        await amfe_cur.for_each(amarker2)
+        assert aaccum2 == [mint(row) for row in base_rows[17:]]
+        assert amfe_cur.state == FindCursorState.CLOSED
+
         # breaking (early) for_each
         accum3: list[dict[str, Any]] = []
 
@@ -354,6 +400,23 @@ def marker3(row: dict[str, Any], acc: list[dict[str, Any]] = accum3) -> bool:
         bfe_another = await bfe_cur.__anext__()
         assert bfe_another == base_rows[5]
 
+        # breaking (early) for_each, coroutine
+        aaccum3: list[dict[str, Any]] = []
+
+        async def amarker3(
+            row: dict[str, Any],
+            acc: list[dict[str, Any]] = aaccum3,
+        ) -> bool:
+            acc += [row]
+            return len(acc) < 5
+
+        abfe_cur = filled_composite_atable.find()
+        await abfe_cur.for_each(amarker3)
+        assert aaccum3 == base_rows[:5]
+        assert abfe_cur.state == FindCursorState.STARTED
+        abfe_another = await abfe_cur.__anext__()
+        assert abfe_another == base_rows[5]
+
         # nonbool-nonbreaking for_each
         accum4: list[dict[str, Any]] = []
 
@@ -366,6 +429,21 @@ def marker4(row: dict[str, Any], acc: list[dict[str, Any]] = accum4) -> int:
         assert accum4 == base_rows
         assert nbfe_cur.state == FindCursorState.CLOSED
 
+        # nonbool-nonbreaking for_each, coroutine
+        aaccum4: list[dict[str, Any]] = []
+
+        async def amarker4(
+            row: dict[str, Any],
+            acc: list[dict[str, Any]] = aaccum4,
+        ) -> int:
+            acc += [row]
+            return 8 if len(acc) < 5 else 0
+
+        anbfe_cur = filled_composite_atable.find()
+        await anbfe_cur.for_each(amarker4)  # type: ignore[arg-type]
+        assert aaccum4 == base_rows
+        assert anbfe_cur.state == FindCursorState.CLOSED
+
     @pytest.mark.describe("test of table cursors, serdes options obeyance, async")
     async def test_table_cursors_serdes_options_async(
         self,