gh-64862: Add the stop_exception parameter in iter() and aiter()

Doc/library/functions.rst

@@ -65,14 +65,54 @@ are always available.  They are listed here in alphabetical order.
 
 
 .. function:: aiter(async_iterable, /)
+              aiter(callable, /, stop_value, *, stop_exception=StopAsyncIteration)
+              aiter(callable, /, *, stop_exception)
+
+   Return an :term:`asynchronous iterator` object.
+   The first argument is interpreted very differently
+   depending on the presence of the other arguments.
+   Without other arguments,
+   the single argument must be an :term:`asynchronous iterable`,
+   and the result is equivalent to calling ``x.__aiter__()``.
+
+   If *stop_value* or *stop_exception* is given,
+   then the first argument must be a callable object.
+   The asynchronous iterator created in this case
+   calls *callable* with no arguments and awaits the result
+   for each call to its :meth:`~object.__anext__` method;
+   if the awaited value is equal to *stop_value*,
+   or if the call raises an exception matching *stop_exception*,
+   :exc:`StopAsyncIteration` will be raised,
+   otherwise the value will be returned.
+   The callable is only called when the result of :meth:`~object.__anext__`
+   is awaited.
+
+   *stop_exception* is an exception class or a tuple of exception classes.
+   If *stop_value* is not specified,
+   the iteration stops only when the callable raises an exception.
+   If the callable raises :exc:`StopAsyncIteration`
+   which does not match *stop_exception*,
+   it is replaced with a :exc:`RuntimeError`,
+   as for asynchronous generators (see :pep:`525`).
+
+   For example, reading fixed-size chunks from an asynchronous stream
+   until the end of file is reached::
 
-   Return an :term:`asynchronous iterator` for an :term:`asynchronous iterable`.
-   Equivalent to calling ``x.__aiter__()``.
+      from functools import partial
+      async for chunk in aiter(partial(reader.read, 1024), b''):
+          process_chunk(chunk)
+
+   Or consuming an :class:`asyncio.Queue` until it is shut down::
 
-   Note: Unlike :func:`iter`, :func:`aiter` has no 2-argument variant.
+      from asyncio import QueueShutDown
+      async for item in aiter(queue.get, stop_exception=QueueShutDown):
+          process_item(item)
 
    .. versionadded:: 3.10
 
+   .. versionchanged:: next
+      Added the *stop_value* and *stop_exception* parameters.
+
 .. function:: all(iterable, /)
 
    Return ``True`` if all elements of the *iterable* are true (or if the iterable
@@ -1143,22 +1183,34 @@ are always available.  They are listed here in alphabetical order.
 
 
 .. function:: iter(iterable, /)
-              iter(callable, sentinel, /)
+              iter(callable, /, stop_value, *, stop_exception=StopIteration)
+              iter(callable, /, *, stop_exception)
 
    Return an :term:`iterator` object.  The first argument is interpreted very
-   differently depending on the presence of the second argument. Without a
-   second argument, the single argument must be a collection object which supports the
+   differently depending on the presence of the other arguments. Without other
+   arguments, the single argument must be a collection object which supports the
    :term:`iterable` protocol (the :meth:`~object.__iter__` method),
    or it must support
    the sequence protocol (the :meth:`~object.__getitem__` method with integer arguments
    starting at ``0``).  If it does not support either of those protocols,
-   :exc:`TypeError` is raised. If the second argument, *sentinel*, is given,
+   :exc:`TypeError` is raised.
+
+   If *stop_value* or *stop_exception* is given,
    then the first argument must be a callable object.  The iterator created in this case
    will call *callable* with no arguments for each call to its
    :meth:`~iterator.__next__` method; if the value returned is equal to
-   *sentinel*, :exc:`StopIteration` will be raised, otherwise the value will
+   *stop_value*, or if the call raises an exception matching *stop_exception*,
+   :exc:`StopIteration` will be raised, otherwise the value will
    be returned.
 
+   *stop_exception* is an exception class or a tuple of exception classes.
+   If *stop_value* is not specified,
+   the iteration stops only when the callable raises an exception.
+   If the callable raises :exc:`StopIteration`
+   which does not match *stop_exception*,
+   it is replaced with a :exc:`RuntimeError`,
+   as for generators (see :pep:`479`).
+
    See also :ref:`typeiter`.
 
    One useful application of the second form of :func:`iter` is to build a
@@ -1170,6 +1222,19 @@ are always available.  They are listed here in alphabetical order.
           for block in iter(partial(f.read, 64), b''):
               process_block(block)
 
+   *stop_exception* is useful for callables
+   which report exhaustion by raising an exception
+   instead of returning a special value.
+   For example, draining a queue::
+
+      import queue
+      for item in iter(input_queue.get_nowait, stop_exception=queue.Empty):
+          process_item(item)
+
+   .. versionchanged:: next
+      Added the *stop_exception* parameter
+      and allowed passing *stop_value* by keyword.
+
 
 .. function:: len(object, /)
 

Doc/whatsnew/3.16.rst

@@ -75,6 +75,13 @@ New features
 Other language changes
 ======================
 
+* The :func:`iter` function now accepts the *stop_exception* parameter.
+  The created iterator stops when the callable raises the specified exception.
+  The second parameter is now named *stop_value* and can be passed by keyword.
+  :func:`aiter` now accepts the same *stop_value* and *stop_exception*
+  parameters, calling an asynchronous callable and awaiting the result.
+  (Contributed by Serhiy Storchaka in :gh:`64862`.)
+
 * :meth:`memoryview.cast` now allows casting a multidimensional
   F-contiguous view to a one-dimensional view.
   (Contributed by Jaemin Park in :gh:`91484`.)

Include/internal/pycore_genobject.h

@@ -29,6 +29,9 @@ PyAPI_FUNC(int) _PyGen_SetStopIterationValue(PyObject *);
 
 // Export for '_asyncio' shared extension
 PyAPI_FUNC(int) _PyGen_FetchStopIterationValue(PyObject **);
+// Set the exception passed to throw(typ[, val[, tb]]).
+// Return 0 on success, -1 on failure.
+extern int _PyGen_SetException(PyObject *typ, PyObject *val, PyObject *tb);
 
 PyAPI_FUNC(PyObject *)_PyCoro_GetAwaitableIter(PyObject *o);
 PyAPI_FUNC(PyObject *)_PyAsyncGenValueWrapperNew(PyThreadState *state, PyObject *);

Include/internal/pycore_global_strings.h

@@ -827,6 +827,8 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(stdout)
         STRUCT_FOR_ID(step)
         STRUCT_FOR_ID(steps)
+        STRUCT_FOR_ID(stop_exception)
+        STRUCT_FOR_ID(stop_value)
         STRUCT_FOR_ID(store_name)
         STRUCT_FOR_ID(strategy)
         STRUCT_FOR_ID(strftime)

Include/internal/pycore_interp_structs.h

@@ -538,7 +538,7 @@ struct _py_func_state {
    If you add a new static type to the standard library, you may have to
    update one of these numbers.
    */
-#define _Py_NUM_MANAGED_PREINITIALIZED_TYPES 120
+#define _Py_NUM_MANAGED_PREINITIALIZED_TYPES 122
 #define _Py_MAX_MANAGED_STATIC_BUILTIN_TYPES \
     (_Py_NUM_MANAGED_PREINITIALIZED_TYPES + 83)
 #define _Py_MAX_MANAGED_STATIC_EXT_TYPES 10

Include/internal/pycore_iterobject.h

@@ -0,0 +1,28 @@
+#ifndef Py_INTERNAL_ITEROBJECT_H
+#define Py_INTERNAL_ITEROBJECT_H
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "this header requires Py_BUILD_CORE define"
+#endif
+
+extern PyTypeObject _PyACallIter_Type;
+extern PyTypeObject _PyACallIterAwaitable_Type;
+
+// Like PyCallIter_New(), but the iteration also stops when *callable* raises
+// an exception matching *stop_exc* (an exception class or a tuple of exception
+// classes).  *sentinel* can be NULL; NULL *stop_exc* means StopIteration.
+extern PyObject *_PyCallIter_NewEx(PyObject *callable, PyObject *sentinel,
+                                   PyObject *stop_exc);
+
+// The asynchronous counterpart of _PyCallIter_NewEx(): the result of
+// *callable* is awaited, and NULL *stop_exc* means StopAsyncIteration.
+extern PyObject *_PyACallIter_New(PyObject *callable, PyObject *sentinel,
+                                  PyObject *stop_exc);
+
+#ifdef __cplusplus
+}
+#endif
+#endif   /* !Py_INTERNAL_ITEROBJECT_H */

Lib/test/test_asyncgen.py

@@ -789,6 +789,164 @@ async def gen():
         applied_twice = aiter(applied_once)
         self.assertIs(applied_once, applied_twice)
 
+    def make_counter(self):
+        state = {'n': 0}
+        async def counter():
+            state['n'] += 1
+            return state['n']
+        return counter
+
+    def collect(self, ait):
+        async def consume():
+            return [i async for i in ait]
+        return self.loop.run_until_complete(consume())
+
+    def test_aiter_callable_stop(self):
+        self.assertEqual(self.collect(aiter(self.make_counter(), 4)), [1, 2, 3])
+        self.assertEqual(self.collect(aiter(self.make_counter(), stop_value=4)),
+                         [1, 2, 3])
+
+    def test_aiter_callable_stop_exception(self):
+        counter = self.make_counter()
+        async def spam():
+            value = await counter()
+            if value > 3:
+                raise LookupError
+            return value
+        self.assertEqual(self.collect(aiter(spam, stop_exception=LookupError)),
+                         [1, 2, 3])
+        counter = self.make_counter()
+        self.assertEqual(
+            self.collect(aiter(spam, stop_exception=(ZeroDivisionError,
+                                                     LookupError))),
+            [1, 2, 3])
+
+    def test_aiter_callable_stop_and_exception(self):
+        counter = self.make_counter()
+        async def spam():
+            value = await counter()
+            if value > 5:
+                raise LookupError
+            return value
+        self.assertEqual(
+            self.collect(aiter(spam, 3, stop_exception=LookupError)), [1, 2])
+        counter = self.make_counter()
+        self.assertEqual(
+            self.collect(aiter(spam, 100, stop_exception=LookupError)),
+            [1, 2, 3, 4, 5])
+
+    def test_aiter_callable_stop_async_iteration(self):
+        # StopAsyncIteration is the default stop exception
+        counter = self.make_counter()
+        async def spam():
+            value = await counter()
+            if value > 3:
+                raise StopAsyncIteration
+            return value
+        self.assertEqual(
+            self.collect(aiter(spam, stop_exception=StopAsyncIteration)),
+            [1, 2, 3])
+
+    def test_aiter_callable_leak_from_await(self):
+        # A StopAsyncIteration leaking from the await is replaced with
+        # RuntimeError (see PEP 525)
+        async def spam():
+            raise StopAsyncIteration
+        it = aiter(spam, 10, stop_exception=LookupError)
+        with self.assertRaisesRegex(RuntimeError,
+                                    'callable raised StopAsyncIteration') as cm:
+            self.loop.run_until_complete(anext(it))
+        self.assertIsInstance(cm.exception.__cause__, StopAsyncIteration)
+        # but if it matches stop_exception, it stops the iteration
+        it = aiter(spam, 10, stop_exception=(LookupError, StopAsyncIteration))
+        with self.assertRaises(StopAsyncIteration):
+            self.loop.run_until_complete(anext(it))
+
+    def test_aiter_callable_leak_from_call(self):
+        # StopIteration and StopAsyncIteration leaking from the call are
+        # replaced with RuntimeError (see PEP 525)
+        for exc in StopIteration, StopAsyncIteration:
+            with self.subTest(exc=exc):
+                def spam():
+                    raise exc
+                it = aiter(spam, 10, stop_exception=LookupError)
+                with self.assertRaisesRegex(
+                        RuntimeError, f'callable raised {exc.__name__}') as cm:
+                    self.loop.run_until_complete(anext(it))
+                self.assertIsInstance(cm.exception.__cause__, exc)
+                # but if it matches stop_exception, it stops the iteration
+                it = aiter(spam, 10, stop_exception=(LookupError, exc))
+                with self.assertRaises(StopAsyncIteration):
+                    self.loop.run_until_complete(anext(it))
+
+    def test_aiter_callable_other_exception(self):
+        async def spam():
+            raise ZeroDivisionError
+        it = aiter(spam, stop_exception=LookupError)
+        with self.assertRaises(ZeroDivisionError):
+            self.loop.run_until_complete(anext(it))
+
+    def test_aiter_callable_exhausted(self):
+        it = aiter(self.make_counter(), 3)
+        self.assertEqual(self.collect(it), [1, 2])
+        self.assertEqual(self.loop.run_until_complete(anext(it, 'default')),
+                         'default')
+        with self.assertRaises(StopAsyncIteration):
+            self.loop.run_until_complete(anext(it))
+
+    def test_aiter_callable_lazy(self):
+        # The callable is only called when the awaitable is awaited
+        calls = []
+        async def spam():
+            calls.append(1)
+            return len(calls)
+        it = aiter(spam, 10)
+        awaitable = it.__anext__()
+        self.assertEqual(calls, [])
+        self.assertEqual(self.loop.run_until_complete(awaitable), 1)
+        self.assertEqual(calls, [1])
+
+    def test_aiter_callable_awaitable(self):
+        it = aiter(self.make_counter(), 10)
+        awaitable = it.__anext__()
+        self.assertIsNone(awaitable.close())
+        with self.assertRaises(RuntimeError):
+            self.loop.run_until_complete(awaitable)
+        awaitable = it.__anext__()
+        with self.assertRaises(KeyError):
+            awaitable.throw(KeyError('injected'))
+
+    def test_aiter_callable_cancel(self):
+        # Cancellation is delivered to the awaited callable result
+        cancelled = []
+        async def spam():
+            try:
+                await asyncio.sleep(10)
+            except asyncio.CancelledError:
+                cancelled.append(1)
+                raise
+        async def consume():
+            async for _ in aiter(spam, None):
+                pass
+        async def main():
+            task = asyncio.ensure_future(consume())
+            await asyncio.sleep(0)
+            task.cancel()
+            with self.assertRaises(asyncio.CancelledError):
+                await task
+        self.loop.run_until_complete(main())
+        self.assertEqual(cancelled, [1])
+
+    def test_aiter_callable_errors(self):
+        async def gen():
+            yield 1
+        self.assertRaises(TypeError, aiter, gen(), 1)
+        self.assertRaises(TypeError, aiter, [1, 2], stop_exception=LookupError)
+        self.assertRaises(TypeError, aiter, len, stop_exception=42)
+        self.assertRaises(TypeError, aiter, len,
+                          stop_exception=(LookupError, 42))
+        self.assertRaises(TypeError, aiter, len, stop_exception=LookupError())
+
     def test_anext_bad_args(self):
         async def gen():
             yield 1