fix AioPageIterator search method (#831)

aiobotocore/paginate.py

@@ -3,6 +3,7 @@
 from botocore.utils import set_value_from_jmespath, merge_dicts
 from botocore.compat import six
 
+import jmespath
 import aioitertools
 
 
@@ -129,6 +130,16 @@ async def build_full_result(self):
             complete_result['NextToken'] = self.resume_token
         return complete_result
 
+    async def search(self, expression):
+        compiled = jmespath.compile(expression)
+        async for page in self:
+            results = compiled.search(page)
+            if isinstance(results, list):
+                for element in results:
+                    yield element
+            else:
+                yield results
+
 
 class AioPaginator(Paginator):
     PAGE_ITERATOR_CLS = AioPageIterator

tests/test_basic_s3.py

@@ -107,6 +107,22 @@ async def test_can_paginate_with_page_size(
     assert key_names == ['key0', 'key1', 'key2', 'key3', 'key4']
 
 
+@pytest.mark.asyncio
+@pytest.mark.moto
+async def test_can_search_paginate(
+        s3_client, bucket_name, create_object):
+    keys = []
+    for i in range(5):
+        key_name = 'key%s' % i
+        keys.append(key_name)
+        await create_object(key_name)
+
+    paginator = s3_client.get_paginator('list_objects')
+    page_iter = paginator.paginate(Bucket=bucket_name)
+    async for key_name in page_iter.search('Contents[*].Key'):
+        assert key_name in keys
+
+
 @pytest.mark.asyncio
 @pytest.mark.moto
 async def test_can_paginate_iterator(s3_client, bucket_name, create_object):