diff --git a/scrapy/utils/python.py b/scrapy/utils/python.py index 5703fd4c3..0bf9bff70 100644 --- a/scrapy/utils/python.py +++ b/scrapy/utils/python.py @@ -12,6 +12,7 @@ from functools import partial, wraps from itertools import chain from scrapy.exceptions import ScrapyDeprecationWarning +from scrapy.utils.asyncgen import as_async_generator from scrapy.utils.decorators import deprecated @@ -355,3 +356,27 @@ class MutableChain: @deprecated("scrapy.utils.python.MutableChain.__next__") def next(self): return self.__next__() + + +async def _async_chain(*iterables): + for it in iterables: + async for o in as_async_generator(it): + yield o + + +class MutableAsyncChain: + """ + Similar to MutableChain but for async iterables + """ + + def __init__(self, *args): + self.data = _async_chain(*args) + + def extend(self, *aiterables): + self.data = _async_chain(self.data, _async_chain(*aiterables)) + + def __aiter__(self): + return self + + async def __anext__(self): + return await self.data.__anext__() diff --git a/tests/test_utils_python.py b/tests/test_utils_python.py index 3115cc92f..58b384591 100644 --- a/tests/test_utils_python.py +++ b/tests/test_utils_python.py @@ -2,15 +2,18 @@ import functools import gc import operator import platform -import unittest from datetime import datetime from itertools import count from warnings import catch_warnings +from twisted.trial import unittest + +from scrapy.utils.asyncgen import as_async_generator, collect_asyncgen +from scrapy.utils.defer import deferred_f_from_coro_f, aiter_errback from scrapy.utils.python import ( memoizemethod_noargs, binary_is_text, equal_attributes, WeakKeyCache, get_func_args, to_bytes, to_unicode, - without_none_values, MutableChain) + without_none_values, MutableChain, MutableAsyncChain) __doctests__ = ['scrapy.utils.python'] @@ -32,6 +35,57 @@ class MutableChainTest(unittest.TestCase): self.assertEqual(list(m), list(range(3, 13))) +class MutableAsyncChainTest(unittest.TestCase): + @staticmethod + async def g1(): + for i in range(3): + yield i + + @staticmethod + async def g2(): + return + yield + + @staticmethod + async def g3(): + for i in range(7, 10): + yield i + + @staticmethod + async def g4(): + for i in range(3, 5): + yield i + 1 / 0 + for i in range(5, 7): + yield i + + @staticmethod + async def collect_asyncgen_exc(asyncgen): + results = [] + async for x in asyncgen: + results.append(x) + return results + + @deferred_f_from_coro_f + async def test_mutableasyncchain(self): + m = MutableAsyncChain(self.g1(), as_async_generator(range(3, 7))) + m.extend(self.g2()) + m.extend(self.g3()) + + self.assertEqual(await m.__anext__(), 0) + results = await collect_asyncgen(m) + self.assertEqual(results, list(range(1, 10))) + + @deferred_f_from_coro_f + async def test_mutableasyncchain_exc(self): + m = MutableAsyncChain(self.g1()) + m.extend(self.g4()) + m.extend(self.g3()) + + results = await collect_asyncgen(aiter_errback(m, lambda _: None)) + self.assertEqual(results, list(range(5))) + + class ToUnicodeTest(unittest.TestCase): def test_converting_an_utf8_encoded_string_to_unicode(self): self.assertEqual(to_unicode(b'lel\xc3\xb1e'), 'lel\xf1e')