Skip to content

Commit

Permalink
Add bot-initiated messaging ability
Browse files Browse the repository at this point in the history
The approach is to use an idle handler, which gets
called whenever something else hasn't already
happened, or once per second. This should help
reduce rate limiting.

Functional testing has been added.

Updated README.md with how to use the new
functionality.

Resolves scrapinghub#92.
  • Loading branch information
csaftoiu authored and oannam committed May 16, 2018
1 parent 4284326 commit 58c9d85
Show file tree
Hide file tree
Showing 9 changed files with 256 additions and 14 deletions.
2 changes: 2 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -14,3 +14,5 @@ slackbot_test_settings.py
/*.egg-info
.cache
/.vscode/
.idea

27 changes: 22 additions & 5 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -111,14 +111,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/private channel chat)
- A function decorated with `listen_to` is called when a message matching the pattern is sent on a channel/private channel chat (not directly sent to the bot)
- 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)
- 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)
Expand All @@ -141,6 +141,23 @@ def help(message):

# Start a thread on the original message
message.reply("Here's a threaded reply", in_thread=True)

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:
Expand Down
12 changes: 12 additions & 0 deletions slackbot/bot.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand Down Expand Up @@ -67,6 +68,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):
"""
Expand Down
36 changes: 33 additions & 3 deletions slackbot/dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,21 +135,51 @@ 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:
event_type = event.get('type')
if event_type == 'message':
self._on_new_message(event)
elif event_type in ['channel_created', 'channel_rename',
'group_joined', 'group_rename',
'im_created']:
'group_joined', 'group_rename',
'im_created']:
channel = [event['channel']]
self._client.parse_channel_data(channel)
elif event_type in ['team_join', 'user_change']:
user = [event['user']]
self._client.parse_user_data(user)
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
Expand Down
7 changes: 6 additions & 1 deletion slackbot/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'):
Expand Down Expand Up @@ -74,3 +75,7 @@ def get_plugins(self, category, text):

if not has_matching_plugin:
yield None, None

def get_idle_plugins(self):
for c in self.idle_commands:
yield c
51 changes: 48 additions & 3 deletions slackbot/plugins/hello.py
Original file line number Diff line number Diff line change
@@ -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)
Expand Down Expand Up @@ -57,3 +58,47 @@ def hello_unicode_message(message):
@listen_to('start a thread')
def start_thread(message):
message.reply('I started a thread', in_thread=True)

# 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")

73 changes: 71 additions & 2 deletions slackbot/slackclient.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,12 +30,16 @@ def __init__(self, token, timeout=None, bot_icon=None, bot_emoji=None, connect=T
self.websocket = None
self.users = {}
self.channels = {}
self.dm_channels = {} # map user id to direct message channel id
self.connected = False
if timeout is None:
self.webapi = slacker.Slacker(self.token)
else:
self.webapi = slacker.Slacker(self.token, timeout=timeout)

# keep track of last action for idle handling
self._last_action = time.time()

if connect:
self.rtm_connect()

Expand Down Expand Up @@ -76,17 +80,25 @@ 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 parse_user_data(self, user_data):
self.users.update({u['id']: u for u in user_data})

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 """
Expand Down Expand Up @@ -116,6 +128,7 @@ def rtm_read(self):
return data

def rtm_send_message(self, channel, message, attachments=None, thread_ts=None):
channel = self._channelify(channel)
message_json = {
'type': 'message',
'channel': channel,
Expand All @@ -126,11 +139,13 @@ def rtm_send_message(self, channel, message, attachments=None, thread_ts=None):
self.send_to_websocket(message_json)

def upload_file(self, channel, fname, fpath, comment):
channel = self._channelify(channel)
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 upload_content(self, channel, fname, content, comment):
self.webapi.files.upload(None,
Expand All @@ -140,6 +155,7 @@ def upload_content(self, channel, fname, content, comment):
initial_comment=comment)

def send_message(self, channel, message, attachments=None, as_user=True, thread_ts=None):
channel = self._channelify(channel)
self.webapi.chat.post_message(
channel,
message,
Expand All @@ -149,10 +165,57 @@ def send_message(self, channel, message, attachments=None, as_user=True, thread_
attachments=attachments,
as_user=as_user,
thread_ts=thread_ts)
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:
Expand All @@ -175,6 +238,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):
Expand Down
17 changes: 17 additions & 0 deletions slackbot/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,3 +77,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
Loading

0 comments on commit 58c9d85

Please sign in to comment.