diff --git a/channel_util.py b/channel_util.py index efddc35..fd073f6 100644 --- a/channel_util.py +++ b/channel_util.py @@ -13,7 +13,7 @@ HOUSEJOBS = "CDWDDTAT0" # Callback for telling what channel we in -def channel_check_callback(slack: SlackClient, msg: dict, match: Match) -> None: +async def channel_check_callback(slack: SlackClient, msg: dict, match: Match) -> None: # Sets the users scroll rest_of_msg = match.group(1).strip() rest_of_msg = rest_of_msg.replace("<", "lcaret") diff --git a/identifier.py b/identifier.py index ebf7fb3..cd55087 100644 --- a/identifier.py +++ b/identifier.py @@ -19,7 +19,7 @@ NON_REG_MSG = ("You currently have no scroll registered. To register, type\n" "except with your scroll instead of 666") -def identify_callback(slack, msg, match): +async def identify_callback(slack, msg, match): """ Sets the users scroll """ @@ -39,7 +39,7 @@ def identify_callback(slack, msg, match): slack_util.reply(slack, msg, result) -def identify_other_callback(slack: SlackClient, msg: dict, match: Match): +async def identify_other_callback(slack: SlackClient, msg: dict, match: Match): """ Sets another users scroll """ @@ -63,7 +63,7 @@ def identify_other_callback(slack: SlackClient, msg: dict, match: Match): # noinspection PyUnusedLocal -def check_callback(slack: SlackClient, msg: dict, match: Match): +async def check_callback(slack: SlackClient, msg: dict, match: Match): """ Replies with the users current scroll assignment """ @@ -78,7 +78,7 @@ def check_callback(slack: SlackClient, msg: dict, match: Match): # noinspection PyUnusedLocal -def name_callback(slack, msg, match): +async def name_callback(slack, msg, match): """ Tells the user what it thinks the calling users name is. """ @@ -119,7 +119,7 @@ def lookup_slackid_brother(slack_id: str) -> Optional[scroll_util.Brother]: return None -def lookup_brother_userids(brother: scroll_util.Brother) -> List[str]: +async def lookup_brother_userids(brother: scroll_util.Brother) -> List[str]: """ Returns a list of all userids associated with the given brother. diff --git a/job_nagger.py b/job_nagger.py index c3d8a6c..df789e3 100644 --- a/job_nagger.py +++ b/job_nagger.py @@ -52,7 +52,7 @@ def get_jobs(day=None): return jobs -def nag_callback(slack, msg, match): +async def nag_callback(slack, msg, match): # Get the day day = match.group(1).lower().strip() jobs = get_jobs(day) diff --git a/job_signoff.py b/job_signoff.py index 0d0589b..78984c0 100644 --- a/job_signoff.py +++ b/job_signoff.py @@ -78,7 +78,7 @@ def alert_user(slack: SlackClient, name: str, saywhat: str) -> None: print("Warning: unable to find dm for brother {}".format(brother_dict)) -def signoff_callback(slack: SlackClient, msg: dict, match: Match) -> None: +async def signoff_callback(slack: SlackClient, msg: dict, match: Match) -> None: """ Callback to signoff a user. """ @@ -101,7 +101,7 @@ def signoff_callback(slack: SlackClient, msg: dict, match: Match) -> None: slack_util.reply(slack, msg, e.as_response()) -def punish_callback(slack: SlackClient, msg: dict, match: Match) -> None: +async def punish_callback(slack: SlackClient, msg: dict, match: Match) -> None: """ Undoes a signoff. Maybe should rename """ @@ -132,7 +132,7 @@ def punish_callback(slack: SlackClient, msg: dict, match: Match) -> None: # noinspection PyUnusedLocal -def reset_callback(slack: SlackClient, msg: dict, match: Match) -> None: +async def reset_callback(slack: SlackClient, msg: dict, match: Match) -> None: """ Resets the scores. """ diff --git a/main.py b/main.py index 430574b..b3b3019 100644 --- a/main.py +++ b/main.py @@ -1,3 +1,4 @@ +import typing from typing import List from slackclient import SlackClient # Obvious @@ -10,6 +11,7 @@ import scroll_util import slack_util import slavestothemachine import job_signoff +import asyncio from dummy import FakeClient # Read api token from file @@ -20,13 +22,9 @@ api_file.close() # Enable to use dummy DEBUG_MODE = False - def main() -> None: wrap = ClientWrapper() - # DEBUG: Add blanked handling - # wrapper.add_hook(".*", print) - # Add scroll handling wrap.add_hook(scroll_util.scroll_hook) @@ -58,12 +56,16 @@ def main() -> None: help_callback = management_commands.list_hooks_callback_gen(wrap.hooks) wrap.add_hook(slack_util.Hook(help_callback, pattern=management_commands.bot_help_pattern)) - wrap.listen() + event_loop = asyncio.get_event_loop() + event_loop.run_until_complete(wrap.listen()) -# Callback to list command hooks - class ClientWrapper(object): + """ + Essentially the main state object. + We only ever expect one of these. + Holds a slack client, and handles messsages. + """ def __init__(self): # Init slack if DEBUG_MODE: @@ -80,9 +82,9 @@ class ClientWrapper(object): def add_hook(self, hook: slack_util.Hook) -> None: self.hooks.append(hook) - def listen(self) -> None: - feed = slack_util.message_stream(self.slack) - for msg in feed: + async def listen(self) -> None: + feed = self.async_message_feed() + async for msg in feed: print(msg) # We only care about standard messages, not subtypes, as those usually just channel activity @@ -106,13 +108,29 @@ class ClientWrapper(object): success = False for hook in self.hooks: - if hook.check(slack_to_use, msg): + if await hook.check(slack_to_use, msg): success = True break if not success: print("No hit on {}".format(msg['text'])) + async def async_message_feed(self) -> typing.AsyncGenerator[dict, None]: + """ + Async wrapper around the message feed. + Yields messages awaitably forever. + """ + # Create the msg feed + feed = slack_util.message_stream(self.slack) + + # Create a simple callable that gets one message from the feed + def get_one(): + return next(feed) + + # Continuously yield async threaded tasks that poll the feed + while True: + yield await asyncio.get_running_loop().run_in_executor(None, get_one) + # run main if __name__ == '__main__': diff --git a/management_commands.py b/management_commands.py index 73c4287..4da9866 100644 --- a/management_commands.py +++ b/management_commands.py @@ -6,9 +6,9 @@ import channel_util import slack_util -def list_hooks_callback_gen(hooks: List[slack_util.Hook]): +async def list_hooks_callback_gen(hooks: List[slack_util.Hook]): # noinspection PyUnusedLocal - def callback(slack, msg, match): + async def callback(slack, msg, match): slack_util.reply(slack, msg, "\n".join(hook.pattern for hook in hooks)) return callback @@ -16,7 +16,7 @@ def list_hooks_callback_gen(hooks: List[slack_util.Hook]): # Gracefully reboot to reload code changes # noinspection PyUnusedLocal -def reboot_callback(slack: SlackClient, msg: dict, match: Match) -> None: +async def reboot_callback(slack: SlackClient, msg: dict, match: Match) -> None: response = "Ok. Rebooting..." slack_util.reply(slack, msg, response) exit(0) diff --git a/scroll_util.py b/scroll_util.py index 343de5a..830c0be 100644 --- a/scroll_util.py +++ b/scroll_util.py @@ -31,7 +31,7 @@ brothers_matches = [m for m in brothers_matches if m] brothers: List[Brother] = [Brother(m.group(2), int(m.group(1))) for m in brothers_matches] -def scroll_callback(slack: SlackClient, msg: dict, match: Match) -> None: +async def scroll_callback(slack: SlackClient, msg: dict, match: Match) -> None: """ Finds the scroll of a brother, or the brother of a scroll, based on msg text. """ diff --git a/slack_util.py b/slack_util.py index b871b25..c93d427 100644 --- a/slack_util.py +++ b/slack_util.py @@ -4,7 +4,7 @@ import re from slackclient import SlackClient import channel_util -from typing import Any, Optional, Generator, Match, Callable, List +from typing import Any, Optional, Generator, Match, Callable, List, Awaitable """ Slack helpers. Separated for compartmentalization @@ -61,7 +61,7 @@ class SlackDebugCondom(object): return self.actual_slack.__getattribute__(name) -def message_stream(slack) -> Generator[dict, None, None]: +def message_stream(slack: SlackClient) -> Generator[dict, None, None]: """ Generator that yields messages from slack. Messages are in standard api format, look it up. @@ -72,19 +72,19 @@ def message_stream(slack) -> Generator[dict, None, None]: if slack.rtm_connect(with_team_state=False, auto_reconnect=True): print("Waiting for messages") while True: - sleep(2) + sleep(1) update = slack.rtm_read() for item in update: if item.get('type') == 'message': yield item - sleep(15) + sleep(5) print("Connection failed - retrying") class Hook(object): def __init__(self, - callback: Callable[[SlackClient, dict, Match], None], + callback: Callable[[SlackClient, dict, Match], Awaitable[None]], # TODO: Fix this type pattern: str = None, channel_whitelist: Optional[List[str]] = None, channel_blacklist: Optional[List[str]] = None): @@ -106,7 +106,7 @@ class Hook(object): else: raise Exception("Cannot whitelist and blacklist") - def check(self, slack: SlackClient, msg: dict) -> bool: + async def check(self, slack: SlackClient, msg: dict) -> bool: # Fail if pattern invalid match = re.match(self.pattern, msg['text'], flags=re.IGNORECASE) if match is None: @@ -125,5 +125,5 @@ class Hook(object): # Otherwise do callback and return success print("Matched on pattern {} callback {}".format(self.pattern, self.callback)) - self.callback(slack, msg, match) + await self.callback(slack, msg, match) return True diff --git a/slavestothemachine.py b/slavestothemachine.py index 0990401..d7b3ac8 100644 --- a/slavestothemachine.py +++ b/slavestothemachine.py @@ -19,7 +19,7 @@ def fmt_work_dict(work_dict: dict) -> str: # noinspection PyUnusedLocal -def count_work_callback(slack: SlackClient, msg: dict, match: Match) -> None: +async def count_work_callback(slack: SlackClient, msg: dict, match: Match) -> None: with shelve.open(DB_NAME) as db: text = msg["text"].lower().strip() @@ -64,7 +64,7 @@ def count_work_callback(slack: SlackClient, msg: dict, match: Match) -> None: # noinspection PyUnusedLocal -def dump_work_callback(slack: SlackClient, msg: dict, match: Match) -> None: +async def dump_work_callback(slack: SlackClient, msg: dict, match: Match) -> None: with shelve.open(DB_NAME) as db: # Dump out each user keys = db.keys()