From b8389b78c1eb637af79cb7035dce6a17ec41e739 Mon Sep 17 00:00:00 2001
From: PIG208 <359101898@qq.com>
Date: Sat, 15 May 2021 21:03:39 +0800
Subject: [PATCH] zulip-bots: Use context manager for incrementor.

---
 .../bots/incrementor/incrementor.py           | 38 +++++++++----------
 1 file changed, 19 insertions(+), 19 deletions(-)

diff --git a/zulip_bots/zulip_bots/bots/incrementor/incrementor.py b/zulip_bots/zulip_bots/bots/incrementor/incrementor.py
index 2b860c2..2346316 100644
--- a/zulip_bots/zulip_bots/bots/incrementor/incrementor.py
+++ b/zulip_bots/zulip_bots/bots/incrementor/incrementor.py
@@ -1,7 +1,7 @@
 # See readme.md for instructions on running this code.
 
 from typing import Dict
-from zulip_bots.lib import BotHandler
+from zulip_bots.lib import BotHandler, use_storage
 
 class IncrementorHandler:
     META = {
@@ -24,28 +24,28 @@ class IncrementorHandler:
             storage.put('message_id', None)
 
     def handle_message(self, message: Dict[str, str], bot_handler: BotHandler) -> None:
-        storage = bot_handler.storage
-        num = storage.get('number')
+        with use_storage(bot_handler.storage, ['number']) as storage:
+            num = storage.get("number")
 
-        # num should already be an int, but we do `int()` to force an
-        # explicit type check
-        num = int(num) + 1
+            # num should already be an int, but we do `int()` to force an
+            # explicit type check
+            num = int(num) + 1
 
-        storage.put('number', num)
-        if storage.get('message_id') is not None:
-            result = bot_handler.update_message(dict(
-                message_id=storage.get('message_id'),
-                content=str(num)
-            ))
+            storage.put('number', num)
+            if storage.get('message_id') is not None:
+                result = bot_handler.update_message(dict(
+                    message_id=storage.get('message_id'),
+                    content=str(num)
+                ))
 
-            # When there isn't an error while updating the message, we won't
-            # attempt to send the it again.
-            if result is None or result.get('result') != 'error':
-                return
+                # When there isn't an error while updating the message, we won't
+                # attempt to send the it again.
+                if result is None or result.get('result') != 'error':
+                    return
 
-        message_info = bot_handler.send_reply(message, str(num))
-        if message_info is not None:
-            storage.put('message_id', message_info['id'])
+            message_info = bot_handler.send_reply(message, str(num))
+            if message_info is not None:
+                storage.put('message_id', message_info['id'])
 
 
 handler_class = IncrementorHandler