Add context manager support to producer & consumers
diff --git a/rocketmq/client.py b/rocketmq/client.py
index d5f3a5d..9500099 100644
--- a/rocketmq/client.py
+++ b/rocketmq/client.py
@@ -111,6 +111,12 @@
if self._handle is not None:
ffi_check(dll.DestroyProducer(self._handle))
+ def __enter__(self):
+ self.start()
+
+ def __exit__(self, type, value, traceback):
+ self.shutdown()
+
def send_sync(self, msg):
cres = _CSendResult()
ffi_check(dll.SendMessageSync(self._handle, msg, ctypes.pointer(cres)))
@@ -167,6 +173,12 @@
if self._handle is not None:
ffi_check(dll.DestroyPushConsumer(self._handle))
+ def __enter__(self):
+ self.start()
+
+ def __exit__(self, type, value, traceback):
+ self.shutdown()
+
def set_message_model(self, model):
ffi_check(dll.SetPushConsumerMessageModel(self._handle, model))
@@ -244,6 +256,12 @@
if self._handle is not None:
ffi_check(dll.DestroyPullConsumer(self._handle))
+ def __enter__(self):
+ self.start()
+
+ def __exit__(self, type, value, traceback):
+ self.shutdown()
+
def start(self):
ffi_check(dll.StartPullConsumer(self._handle))
