Merge pull request #17 from ekerstens/context-manager-support

Add context manager support
robinhood · Feb 25, 2020 · 10e4119 · 10e4119
2 parents 96ea83d + b84fe6e
commit 10e4119
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/aiokafka/consumer/consumer.py b/aiokafka/consumer/consumer.py
@@ -321,6 +321,13 @@ def __del__(self, _warnings=warnings):
                 context['source_traceback'] = self._source_traceback
             self._loop.call_exception_handler(context)
 
+    async def __aenter__(self):
+        await self.start()
+        return self
+
+    async def __aexit__(self, type, value, traceback):
+        await self.stop()
+
     async def start(self):
         """ Connect to Kafka cluster. This will:
 

diff --git a/aiokafka/producer/producer.py b/aiokafka/producer/producer.py
@@ -278,6 +278,13 @@ def __del__(self, _warnings=warnings):
                 context['source_traceback'] = self._source_traceback
             self._loop.call_exception_handler(context)
 
+    async def __aenter__(self):
+        await self.start()
+        return self
+
+    async def __aexit__(self, type, value, traceback):
+        await self.stop()
+
     async def start(self):
         """Connect to Kafka cluster and check server version"""
         log.debug("Starting the Kafka producer")  # trace