zulip_bots: Update StateHandler API to behave dict-like.

This matches the external StateHandler API with the embedded
StateHandler API.

zulip_bots/zulip_bots/bots/incrementor/incrementor.py

@@ -11,17 +11,24 @@ class IncrementorHandler(object):
         is @-mentioned, this number will be incremented in the same message.
         '''
 
+    def initialize(self, bot_handler):
+        storage = bot_handler.storage
+        if not storage.contains('number') or not storage.contains('message_id'):
+            storage.put('number', 0)
+            storage.put('message_id', None)
+
     def handle_message(self, message, bot_handler):
-        with bot_handler.storage.state({'number': 0, 'message_id': None}) as state:
+        storage = bot_handler.storage
-            state['number'] += 1
+        num = storage.get('number')
-            if state['message_id'] is None:
+        storage.put('number', num + 1)
-                result = bot_handler.send_reply(message, str(state['number']))
+        if storage.get('message_id') is None:
-                state['message_id'] = result['id']
+            result = bot_handler.send_reply(message, str(storage.get('number')))
-            else:
+            storage.put('message_id', result['id'])
-                bot_handler.update_message(dict(
+        else:
-                    message_id = state['message_id'],
+            bot_handler.update_message(dict(
-                    content = str(state['number'])
+                message_id = storage.get('message_id'),
-                ))
+                content = str(storage.get('number'))
+            ))
 
 
 handler_class = IncrementorHandler

zulip_bots/zulip_bots/bots/incrementor/test_incrementor.py

@@ -11,6 +11,7 @@ class TestIncrementorBot(BotTestCase):
     bot_name = "incrementor"
 
     def test_bot(self):
+        self.initialize_bot()
         messages = [  # Template for message inputs to test, absent of message content
             {
                 'type': 'stream',

zulip_bots/zulip_bots/bots/tictactoe/tictactoe.py

@@ -280,33 +280,33 @@ class ticTacToeHandler(object):
         for val in command_list:
             command += val
         original_sender = message['sender_email']
+        storage = bot_handler.storage
+        if not storage.contains(original_sender):
+            storage.put(original_sender, None)
+        user_board = storage.get(original_sender)
+        if (not user_board) and command == "new":
+            user_board = copy.deepcopy(initial_board)
+            storage.put(original_sender, user_board)
+        user_game = TicTacToeGame(user_board) if user_board else None
 
-        with bot_handler.storage.state({}) as mydict:
+        if command == 'new':
-            user_board = mydict.get(original_sender)
+            if user_game and not first_time(user_game.board):
-            if (not user_board) and command == "new":
+                return_content = "You're already playing a game! Type **@tictactoe help** or **@ttt help** to see valid inputs."
-                user_board = copy.deepcopy(initial_board)
-                mydict[original_sender] = user_board
-            user_game = TicTacToeGame(user_board) if user_board else None
-
-            if command == 'new':
-                if user_game and not first_time(user_game.board):
-                    return_content = "You're already playing a game! Type **@tictactoe help** or **@ttt help** to see valid inputs."
-                else:
-                    return_content = "Welcome to tic-tac-toe! You'll be x's and I'll be o's. Your move first!\n"
-                    return_content += TicTacToeGame.positions
-            elif command == 'help':
-                return_content = TicTacToeGame.detailed_help_message
-            elif (user_game) and TicTacToeGame.check_validity(user_game, TicTacToeGame.sanitize_move(user_game, command)):
-                user_board = user_game.board
-                return_content = TicTacToeGame.tictactoe(user_game, user_board, command)
-            elif (user_game) and command == 'quit':
-                del mydict[original_sender]
-                return_content = "You've successfully quit the game."
             else:
-                return_content = "Hmm, I didn't understand your input. Type **@tictactoe help** or **@ttt help** to see valid inputs."
+                return_content = "Welcome to tic-tac-toe! You'll be x's and I'll be o's. Your move first!\n"
+                return_content += TicTacToeGame.positions
+        elif command == 'help':
+            return_content = TicTacToeGame.detailed_help_message
+        elif (user_game) and TicTacToeGame.check_validity(user_game, TicTacToeGame.sanitize_move(user_game, command)):
+            return_content = TicTacToeGame.tictactoe(user_game, user_board, command)
+        elif (user_game) and command == 'quit':
+            storage.put(original_sender, None)
+            return_content = "You've successfully quit the game."
+        else:
+            return_content = "Hmm, I didn't understand your input. Type **@tictactoe help** or **@ttt help** to see valid inputs."
 
-            if "Game over" in return_content or "draw" in return_content:
+        if "Game over" in return_content or "draw" in return_content:
-                del mydict[original_sender]
+            storage.put(original_sender, None)
 
         bot_handler.send_message(dict(
             type = 'private',

zulip_bots/zulip_bots/bots/virtual_fs/virtual_fs.py

@@ -18,16 +18,16 @@ class VirtualFsHandler(object):
         if isinstance(recipient, list):  # If not a stream, then hash on list of emails
             recipient = " ".join([x['email'] for x in recipient])
 
-        with bot_handler.storage.state({}) as state:
+        storage = bot_handler.storage
-            if recipient not in state:
+        if not storage.contains(recipient):
-                state[recipient] = fs_new()
+            storage.put(recipient, fs_new())
-            fs = state[recipient]
+        fs = storage.get(recipient)
-            if sender not in fs['user_paths']:
+        if sender not in fs['user_paths']:
-                fs['user_paths'][sender] = '/'
+            fs['user_paths'][sender] = '/'
-            fs, msg = fs_command(fs, sender, command)
+        fs, msg = fs_command(fs, sender, command)
-            prependix = '{}:\n'.format(sender)
+        prependix = '{}:\n'.format(sender)
-            msg = prependix + msg
+        msg = prependix + msg
-            state[recipient] = fs
+        storage.put(recipient, fs)
 
         bot_handler.send_reply(message, msg)
 

zulip_bots/zulip_bots/lib.py

@@ -1,5 +1,6 @@
 from __future__ import print_function
 
+import json
 import logging
 import os
 import signal
@@ -54,22 +55,21 @@ class RateLimit(object):
 class StateHandler(object):
     def __init__(self):
         # type: () -> None
-        self.state_ = None  # type: Any
+        self.state_ = {}  # type: Dict[Text, Text]
+        self.marshal = lambda obj: obj
+        self.demarshal = lambda obj: obj
 
-    def set_state(self, state):
+    def put(self, key, value):
-        # type: (Any) -> None
+        # type: (Text, Text) -> None
-        self.state_ = state
+        self.state_[key] = self.marshal(value)
 
-    def get_state(self):
+    def get(self, key):
-        # type: () -> Any
+        # type: () -> Text
-        return self.state_
+        return self.demarshal(self.state_[key])
 
-    @contextmanager
+    def contains(self, key):
-    def state(self, default):
+        # type: (Text) -> bool
-        # type: (Any) -> Any
+        return key in self.state_
-        new_state = self.get_state() or default
-        yield new_state
-        self.set_state(new_state)
 
 class ExternalBotHandler(object):
     def __init__(self, client, root_dir):