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