Implement async contest manager interface on AsyncElasticsearch

Features:
- updated README with the new syntax (first example)
- test code in a separate file, because it has to use Python 3.5+ syntax

Author: Albin Stjerna

README

@@ -11,17 +11,14 @@ Example::
 
     import asyncio
     from elasticsearch_async import AsyncElasticsearch
+    hosts = ['localhost', 'other-host']
 
-    client = AsyncElasticsearch(hosts=['localhost', 'other-host'])
-
-    @asyncio.coroutine
-    def print_info():
-        info = yield from client.info()
-        print(info)
+    async def print_info():
+        async with AsyncElasticsearch(hosts=hosts) as client:
+              print(await client.info())
 
     loop = asyncio.get_event_loop()
     loop.run_until_complete(print_info())
-    loop.run_until_complete(client.transport.close())
     loop.close()
 
 

elasticsearch_async/__init__.py

@@ -1,8 +1,19 @@
-from .transport import AsyncTransport
-from .connection import AIOHttpConnection
+import asyncio
 
 from elasticsearch import Elasticsearch
 
+from .connection import AIOHttpConnection
+from .transport import AsyncTransport
+
+
 class AsyncElasticsearch(Elasticsearch):
     def __init__(self, hosts=None, transport_class=AsyncTransport, **kwargs):
         super().__init__(hosts, transport_class=transport_class, **kwargs)
+
+    @asyncio.coroutine
+    def __aenter__(self):
+        return self
+
+    @asyncio.coroutine
+    def __aexit__(self, _exc_type, _exc_val, _exc_tb):
+        yield from self.transport.close()

test_elasticsearch_async/test_transport.py

@@ -4,11 +4,13 @@
 
 from elasticsearch_async import AsyncElasticsearch
 
+
 @mark.asyncio
 def test_sniff_on_start_sniffs(server, event_loop, port, sniff_data):
     server.register_response('/_nodes/_all/http', sniff_data)
 
-    client = AsyncElasticsearch(port=port, sniff_on_start=True, loop=event_loop)
+    client = AsyncElasticsearch(
+        port=port, sniff_on_start=True, loop=event_loop)
 
     # sniff has been called in the background
     assert client.transport.sniffing_task is not None
@@ -21,10 +23,15 @@ def test_sniff_on_start_sniffs(server, event_loop, port, sniff_data):
     assert 'http://node1:9200' == connections[0].host
     yield from client.transport.close()
 
+
 @mark.asyncio
 def test_retry_will_work(port, server, event_loop):
-    client = AsyncElasticsearch(hosts=['not-an-es-host', 'localhost'], port=port, loop=event_loop, randomize_hosts=False)
+    client = AsyncElasticsearch(
+        hosts=['not-an-es-host', 'localhost'],
+        port=port,
+        loop=event_loop,
+        randomize_hosts=False)
 
     data = yield from client.info()
-    assert  {'body': '', 'method': 'GET', 'params': {}, 'path': '/'} == data
+    assert {'body': '', 'method': 'GET', 'params': {}, 'path': '/'} == data
     yield from client.transport.close()

test_elasticsearch_async/test_with_syntax.py

@@ -0,0 +1,15 @@
+import asyncio
+import sys
+
+from pytest import mark
+
+from elasticsearch_async import AsyncElasticsearch
+
+
+@mark.skipif(sys.version_info < (3, 5), reason="async with is Python 3.5+")
+@mark.asyncio
+async def test_with_notation_works(port, server, event_loop):
+    async with AsyncElasticsearch(
+            hosts=['localhost'], port=port, loop=event_loop) as client:
+        info = await client.info()
+    assert info