diff --git a/.gitignore b/.gitignore index d3e1bedc..ca5fb4e9 100644 --- a/.gitignore +++ b/.gitignore @@ -13,3 +13,5 @@ slackbot_test_settings.py /dist /*.egg-info .cache +.idea + diff --git a/README.md b/README.md index 0052ae0a..aac24c6a 100644 --- a/README.md +++ b/README.md @@ -108,14 +108,14 @@ def github(): A chat bot is meaningless unless you can extend/customize it to fit your own use cases. -To write a new plugin, simplely create a function decorated by `slackbot.bot.respond_to` or `slackbot.bot.listen_to`: +To write a new plugin, simply create a function decorated by `slackbot.bot.respond_to`, `slackbot.bot.listen_to`, or `slackbot.bot.idle`: - A function decorated with `respond_to` is called when a message matching the pattern is sent to the bot (direct message or @botname in a channel/group chat) - A function decorated with `listen_to` is called when a message matching the pattern is sent on a channel/group chat (not directly sent to the bot) +- *(development version only)* A function decorated with `idle` is called whenever a message has not been sent for the past second ```python -from slackbot.bot import respond_to -from slackbot.bot import listen_to +from slackbot.bot import respond_to, listen_to, idle import re @respond_to('hi', re.IGNORECASE) @@ -135,6 +135,23 @@ def help(message): # Message is sent on the channel # message.send('I can help everybody!') + +last_bored = time.time() +@idle +def bored(client): + if time.time() - last_bored >= 30: + last_bored = time.time() + + # Messages can be sent to a channel + client.rtm_send_message('some_channel', "I'm bored!") + # Or directly to a user + client.rtm_send_message('some_user', "Hey, entertain me!") + + # If a name is ambiguous: + client.rtm_send_message(client.find_channel_by_name('ambiguous'), "To ambiguous the channel") + client.rtm_send_message(client.find_user_by_name('ambiguous'), "To ambiguous the user") + + # Attachments can be sent with `client.rtm_send_message(..., attachments=attachments)`. ``` To extract params from the message, you can use regular expression: diff --git a/slackbot/bot.py b/slackbot/bot.py index a45010af..39592678 100644 --- a/slackbot/bot.py +++ b/slackbot/bot.py @@ -11,6 +11,7 @@ from slackbot.manager import PluginsManager from slackbot.slackclient import SlackClient from slackbot.dispatcher import MessageDispatcher +from slackbot.utils import optional_arg_decorator logger = logging.getLogger(__name__) @@ -65,6 +66,17 @@ def wrapper(func): return wrapper +# use optional_arg_decorator so users can either do @idle or @idle() +@optional_arg_decorator +def idle(func): + """Run a function once/second whenever no other actions were taken. + The function must take one parameter, a SlackClient instance.""" + # match anything, the text doesn't apply for "idle" + PluginsManager.idle_commands.append(func) + logger.info('registered idle plugin "%s"', func.__name__) + return func + + # def default_reply(matchstr=r'^.*$', flags=0): def default_reply(*args, **kwargs): """ diff --git a/slackbot/dispatcher.py b/slackbot/dispatcher.py index 14f847e7..ff68e1ac 100644 --- a/slackbot/dispatcher.py +++ b/slackbot/dispatcher.py @@ -135,13 +135,43 @@ def filter_text(self, msg): return msg def loop(self): + # once/second, check events + # run idle handlers whenever idle while True: events = self._client.rtm_read() for event in events: if event.get('type') != 'message': continue self._on_new_message(event) - time.sleep(1) + + # run idle handlers as long as we've been idle + for func in self._plugins.get_idle_plugins(): + if not func: + continue + + # if actions are pending, don't do anything + if not self._pool.queue.empty(): + break + + # if some action was taken, don't run the remaining handlers + if self._client.idle_time() < 1: + break + + try: + func(self._client) + except: + logger.exception( + 'idle handler failed with plugin "%s"', + func.__name__) + reply = u'[{}] I had a problem with idle handler\n'.format( + func.__name__) + tb = u'```\n{}\n```'.format(traceback.format_exc()) + # no channel, so only send errors to error user + if self._errors_to: + self._client.rtm_send_message(self._errors_to, + '{}\n{}'.format(reply, + tb)) + time.sleep(1.0) def _default_reply(self, msg): default_reply = settings.DEFAULT_REPLY diff --git a/slackbot/manager.py b/slackbot/manager.py index 0239e238..48859358 100644 --- a/slackbot/manager.py +++ b/slackbot/manager.py @@ -18,8 +18,9 @@ def __init__(self): commands = { 'respond_to': {}, 'listen_to': {}, - 'default_reply': {} + 'default_reply': {}, } + idle_commands = [] def init_plugins(self): if hasattr(settings, 'PLUGINS'): @@ -72,3 +73,6 @@ def get_plugins(self, category, text): if not has_matching_plugin: yield None, None + + def get_idle_plugins(self): + yield from self.idle_commands diff --git a/slackbot/plugins/hello.py b/slackbot/plugins/hello.py index d2b81937..1b7b0893 100644 --- a/slackbot/plugins/hello.py +++ b/slackbot/plugins/hello.py @@ -1,7 +1,8 @@ -#coding: UTF-8 +# coding: UTF-8 +import random import re -from slackbot.bot import respond_to -from slackbot.bot import listen_to + +from slackbot.bot import respond_to, listen_to, idle @respond_to('hello$', re.IGNORECASE) @@ -52,3 +53,47 @@ def hey(message): @respond_to(u'你好') def hello_unicode_message(message): message.reply(u'你好!') + + +# idle tests +IDLE_TEST = {'which': None, 'channel': None} + + +@respond_to('start idle test ([0-9]+)') +@listen_to('start idle test ([0-9]+)') +def start_idle_test(message, i): + print("---------- start idle test! -----------") + IDLE_TEST['which'] = int(i) + IDLE_TEST['channel'] = message._body['channel'] + print("Idle test is now {which} on channel {channel}".format(**IDLE_TEST)) + # TESTING ONLY, don't rely on this behavior + + +# idle function testing +# tests 0 and 1: rtm and webapi work from idle function 1 +# tests 2 and 3: rtm and webapi work from idle function 2 +# test 4: both idle functions can operate simultaneously +@idle +def idle_1(client): + which = IDLE_TEST['which'] + msg = "I am bored %s" % which + if which == 0: + client.rtm_send_message(IDLE_TEST['channel'], msg) + elif which == 1: + client.send_message(IDLE_TEST['channel'], msg) + elif which == 4: + if random.random() <= 0.5: + client.rtm_send_message(IDLE_TEST['channel'], "idle_1 is bored") + + +@idle() +def idle_2(client): + which = IDLE_TEST['which'] + msg = "I am bored %s" % which + if which == 2: + client.rtm_send_message(IDLE_TEST['channel'], msg) + elif which == 3: + client.send_message(IDLE_TEST['channel'], msg) + elif which == 4: + if random.random() <= 0.5: + client.rtm_send_message(IDLE_TEST['channel'], "idle_2 is bored") diff --git a/slackbot/slackclient.py b/slackbot/slackclient.py index 6c10992f..f180d65e 100644 --- a/slackbot/slackclient.py +++ b/slackbot/slackclient.py @@ -30,9 +30,13 @@ def __init__(self, token, bot_icon=None, bot_emoji=None, connect=True): self.websocket = None self.users = {} self.channels = {} + self.dm_channels = {} # map user id to direct message channel id self.connected = False self.webapi = slacker.Slacker(self.token) + # keep track of last action for idle handling + self._last_action = time.time() + if connect: self.rtm_connect() @@ -65,14 +69,22 @@ def parse_slack_login_data(self, login_data): def parse_channel_data(self, channel_data): self.channels.update({c['id']: c for c in channel_data}) + # pre-load direct message channels + for c in channel_data: + if 'user' in c: + self.dm_channels[c['user']] = c['id'] def send_to_websocket(self, data): - """Send (data) directly to the websocket.""" + """Send (data) directly to the websocket. + + Update last action for idle handling.""" data = json.dumps(data) self.websocket.send(data) + self._last_action = time.time() def ping(self): - return self.send_to_websocket({'type': 'ping'}) + self.send_to_websocket({'type': 'ping'}) + self._last_action = time.time() def websocket_safe_read(self): """Returns data if available, otherwise ''. Newlines indicate multiple messages """ @@ -101,7 +113,8 @@ def rtm_read(self): data.append(json.loads(d)) return data - def rtm_send_message(self, channel, message, attachments=None): + def rtm_send_message(self, channelish, message, attachments=None): + channel = self._channelify(channelish) message_json = { 'type': 'message', 'channel': channel, @@ -110,14 +123,17 @@ def rtm_send_message(self, channel, message, attachments=None): } self.send_to_websocket(message_json) - def upload_file(self, channel, fname, fpath, comment): + def upload_file(self, channelish, fname, fpath, comment): + channel = self._channelify(channelish) fname = fname or to_utf8(os.path.basename(fpath)) self.webapi.files.upload(fpath, channels=channel, filename=fname, initial_comment=comment) + self._last_action = time.time() - def send_message(self, channel, message, attachments=None, as_user=True): + def send_message(self, channelish, message, attachments=None, as_user=True): + channel = self._channelify(channelish) self.webapi.chat.post_message( channel, message, @@ -126,10 +142,57 @@ def send_message(self, channel, message, attachments=None, as_user=True): icon_emoji=self.bot_emoji, attachments=attachments, as_user=as_user) + self._last_action = time.time() def get_channel(self, channel_id): return Channel(self, self.channels[channel_id]) + def get_dm_channel(self, user_id): + """Get the direct message channel for the given user id, opening + one if necessary.""" + if user_id not in self.users: + raise ValueError("Expected valid user_id, have no user '%s'" % ( + user_id,)) + + if user_id in self.dm_channels: + return self.dm_channels[user_id] + + # open a new channel + resp = self.webapi.im.open(user_id) + if not resp.body["ok"]: + raise ValueError("Could not open DM channel: %s" % resp.body) + + self.dm_channels[user_id] = resp.body['channel']['id'] + + return self.dm_channels[user_id] + + def _channelify(self, s): + """Turn a string into a channel. + + * Given a channel id, return that same channel id. + * Given a channel name, return the channel id. + * Given a user id, return the direct message channel with that user, + opening a new one if necessary. + * Given a user name, do the same as for a user id. + + Raise a ValueError otherwise.""" + if s in self.channels: + return s + + channel_id = self.find_channel_by_name(s) + if channel_id: + return channel_id + + if s in self.users: + return self.get_dm_channel(s) + + user_id = self.find_user_by_name(s) + if user_id: + return self.get_dm_channel(user_id) + + raise ValueError("Could not turn '%s' into any kind of channel name" % ( + user_id)) + def find_channel_by_name(self, channel_name): for channel_id, channel in iteritems(self.channels): try: @@ -149,6 +212,12 @@ def react_to_message(self, emojiname, channel, timestamp): name=emojiname, channel=channel, timestamp=timestamp) + self._last_action = time.time() + + def idle_time(self): + """Return the time the client has been idle, i.e. the time since + it sent the last message to the server.""" + return time.time() - self._last_action class SlackConnectionError(Exception): diff --git a/slackbot/utils.py b/slackbot/utils.py index b79911c8..26cf206e 100644 --- a/slackbot/utils.py +++ b/slackbot/utils.py @@ -76,3 +76,20 @@ def do_work(self): while True: msg = self.queue.get() self.func(msg) + + +def optional_arg_decorator(fn): + """Allows for easier making of decorators with optional arguments. + + See: http://stackoverflow.com/questions/3888158/python-making-decorators-with-optional-arguments""" + def wrapped_decorator(*args): + if len(args) == 1 and callable(args[0]): + return fn(args[0]) + + else: + def real_decorator(decoratee): + return fn(decoratee, *args) + + return real_decorator + + return wrapped_decorator diff --git a/tests/functional/test_functional.py b/tests/functional/test_functional.py index 1906efda..d2a24a3f 100644 --- a/tests/functional/test_functional.py +++ b/tests/functional/test_functional.py @@ -8,7 +8,9 @@ import os from os.path import join, abspath, dirname, basename import subprocess + import pytest + from tests.functional.driver import Driver from tests.functional.slackbot_settings import ( testbot_apitoken, testbot_username, @@ -192,3 +194,45 @@ def test_bot_reply_with_alias_message(driver): driver.wait_for_bot_channel_message("hello sender!", tosender=True) driver.send_channel_message('!hello', tobot=False, colon=False) driver.wait_for_bot_channel_message("hello sender!", tosender=True) + + +def make_idle_func(use_rtm): + def idle_func_bored(client): + if use_rtm: + client.rtm_send_message(test_channel, "I am bored") + else: + client.send_message(test_channel, "I am bored") + return idle_func_bored + + +# 5 tests, defined in hello.py +# parametrize based on method for fewer test functions +send_methods = { + 'direct': { + 'send': lambda driver, msg: driver.send_direct_message(msg), + 'wait': lambda driver, msg: driver.wait_for_bot_direct_message(msg), + }, + 'channel': { + 'send': lambda driver, msg: driver.send_channel_message(msg), + 'wait': lambda driver, msg: driver.wait_for_bot_channel_message(msg, tosender=False) + }, + 'group': { + 'send': lambda driver, msg: driver.send_group_message(msg), + 'wait': lambda driver, msg: driver.wait_for_bot_group_message(msg, tosender=False) + }, +} + +@pytest.mark.parametrize('which_test', [0, 1, 2, 3]) +@pytest.mark.parametrize('method', send_methods.keys()) +def test_idle_func_dm(driver, which_test, method): + send_methods[method]['send'](driver, "start idle test %d" % which_test) + send_methods[method]['wait'](driver, "I am bored %d" % which_test) + + +@pytest.mark.parametrize('method', send_methods.keys()) +def test_idle_func_dm_multi(driver, method): + send_methods[method]['send'](driver, "start idle test 4") + send_methods[method]['wait'](driver, "idle_1 is bored") + send_methods[method]['wait'](driver, "idle_2 is bored") + send_methods[method]['wait'](driver, "idle_1 is bored") + send_methods[method]['wait'](driver, "idle_2 is bored")