zulip_bots: Update StateHandler API to behave dict-like.

This matches the external StateHandler API with the embedded
StateHandler API.
This commit is contained in:
derAnfaenger 2017-10-24 13:15:51 +02:00
parent 32df4e097d
commit 59f81845dd
5 changed files with 65 additions and 57 deletions

View file

@ -11,16 +11,23 @@ class IncrementorHandler(object):
is @-mentioned, this number will be incremented in the same message. 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): 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')))
storage.put('message_id', result['id'])
else: else:
bot_handler.update_message(dict( bot_handler.update_message(dict(
message_id = state['message_id'], message_id = storage.get('message_id'),
content = str(state['number']) content = str(storage.get('number'))
)) ))

View file

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

View file

@ -280,12 +280,13 @@ class ticTacToeHandler(object):
for val in command_list: for val in command_list:
command += val command += val
original_sender = message['sender_email'] original_sender = message['sender_email']
storage = bot_handler.storage
with bot_handler.storage.state({}) as mydict: if not storage.contains(original_sender):
user_board = mydict.get(original_sender) storage.put(original_sender, None)
user_board = storage.get(original_sender)
if (not user_board) and command == "new": if (not user_board) and command == "new":
user_board = copy.deepcopy(initial_board) user_board = copy.deepcopy(initial_board)
mydict[original_sender] = user_board storage.put(original_sender, user_board)
user_game = TicTacToeGame(user_board) if user_board else None user_game = TicTacToeGame(user_board) if user_board else None
if command == 'new': if command == 'new':
@ -297,16 +298,15 @@ class ticTacToeHandler(object):
elif command == 'help': elif command == 'help':
return_content = TicTacToeGame.detailed_help_message return_content = TicTacToeGame.detailed_help_message
elif (user_game) and TicTacToeGame.check_validity(user_game, TicTacToeGame.sanitize_move(user_game, command)): 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) return_content = TicTacToeGame.tictactoe(user_game, user_board, command)
elif (user_game) and command == 'quit': elif (user_game) and command == 'quit':
del mydict[original_sender] storage.put(original_sender, None)
return_content = "You've successfully quit the game." return_content = "You've successfully quit the game."
else: else:
return_content = "Hmm, I didn't understand your input. Type **@tictactoe help** or **@ttt help** to see valid inputs." 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( bot_handler.send_message(dict(
type = 'private', type = 'private',

View file

@ -18,16 +18,16 @@ class VirtualFsHandler(object):
if isinstance(recipient, list): # If not a stream, then hash on list of emails if isinstance(recipient, list): # If not a stream, then hash on list of emails
recipient = " ".join([x['email'] for x in recipient]) 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) bot_handler.send_reply(message, msg)

View file

@ -1,5 +1,6 @@
from __future__ import print_function from __future__ import print_function
import json
import logging import logging
import os import os
import signal import signal
@ -54,22 +55,21 @@ class RateLimit(object):
class StateHandler(object): class StateHandler(object):
def __init__(self): def __init__(self):
# type: () -> None # 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): class ExternalBotHandler(object):
def __init__(self, client, root_dir): def __init__(self, client, root_dir):