diff --git a/zulip_bots/setup.py b/zulip_bots/setup.py index 207d43c0e..8eff799eb 100644 --- a/zulip_bots/setup.py +++ b/zulip_bots/setup.py @@ -55,6 +55,7 @@ 'html2text', 'lxml', 'BeautifulSoup4', + 'entrypoints', ], ) diff --git a/zulip_bots/zulip_bots/finder.py b/zulip_bots/zulip_bots/finder.py index 0c44b1b65..8e71c9208 100644 --- a/zulip_bots/zulip_bots/finder.py +++ b/zulip_bots/zulip_bots/finder.py @@ -1,5 +1,6 @@ import sys import os +import entrypoints from os.path import basename, splitext from typing import Any, Optional, Text, Tuple @@ -32,6 +33,26 @@ def import_module_by_name(name: Text) -> Any: except ModuleNotFoundError: # Specific exception supported >=Python3.6 return None +class DuplicateRegisteredBotName(Exception): + pass + +def import_module_from_zulip_bot_registry(name: str) -> Tuple[str, Any]: + registered_bots = entrypoints.get_group_all('zulip_bots.registry') + matching_bots = [bot for bot in registered_bots if bot.name == name] + + if len(matching_bots) == 1: # Unique matching entrypoint + bot = matching_bots[0] + if bot.distro is not None: + return "{}: {}".format(bot.distro.name, bot.distro.version), bot.load() + else: + print(bot) + return "editable package: {}".format(bot.module_name), bot.load() + + if len(matching_bots) > 1: + raise DuplicateRegisteredBotName(name) + + return "", None # no matches in registry + def resolve_bot_path(name: Text) -> Optional[Tuple[Text, Text]]: if os.path.isfile(name): bot_path = os.path.abspath(name) diff --git a/zulip_bots/zulip_bots/lib.py b/zulip_bots/zulip_bots/lib.py index 55c7e37af..0798be517 100644 --- a/zulip_bots/zulip_bots/lib.py +++ b/zulip_bots/zulip_bots/lib.py @@ -304,6 +304,7 @@ def run_message_handler_for_bot( config_file: str, bot_config_file: str, bot_name: str, + bot_source: str, ) -> Any: """ lib_module is of type Any, since it can contain any bot's @@ -334,7 +335,7 @@ def run_message_handler_for_bot( message_handler = prepare_message_handler(bot_name, restricted_client, lib_module) if not quiet: - print("Running {} Bot:".format(bot_details['name'])) + print("Running {} Bot (from {}):".format(bot_details['name'], bot_source)) if bot_details['description'] != "": print("\n\t{}".format(bot_details['description'])) print(message_handler.usage()) diff --git a/zulip_bots/zulip_bots/run.py b/zulip_bots/zulip_bots/run.py index 0634bf8fd..27d988d23 100755 --- a/zulip_bots/zulip_bots/run.py +++ b/zulip_bots/zulip_bots/run.py @@ -123,13 +123,23 @@ def main() -> None: " zulip-run-bot {bot_name} --provision") print(dep_err_msg.format(bot_name=bot_name, deps_list=deps_list)) sys.exit(1) + bot_source = "source" else: lib_module = finder.import_module_by_name(args.bot) if lib_module: bot_name = lib_module.__name__ + bot_source = "named module" if args.provision: print("ERROR: Could not load bot's module for '{}'. Exiting now.") sys.exit(1) + else: + try: + bot_source, lib_module = finder.import_module_from_zulip_bot_registry(args.bot) + except finder.DuplicateRegisteredBotName as e: + print("ERROR: Found duplicate entries for bot name in zulip bot registry. Exiting now.") + sys.exit(1) + if lib_module: + bot_name = args.bot if lib_module is None: print("ERROR: Could not load bot module. Exiting now.") @@ -150,7 +160,8 @@ def main() -> None: config_file=args.config_file, bot_config_file=args.bot_config_file, quiet=args.quiet, - bot_name=bot_name + bot_name=bot_name, + bot_source=bot_source, ) except NoBotConfigException: print(''' diff --git a/zulip_bots/zulip_bots/tests/test_lib.py b/zulip_bots/zulip_bots/tests/test_lib.py index dbf882e77..79d2dad02 100644 --- a/zulip_bots/zulip_bots/tests/test_lib.py +++ b/zulip_bots/zulip_bots/tests/test_lib.py @@ -174,7 +174,8 @@ def test_message(message, flags): quiet=True, config_file=None, bot_config_file=None, - bot_name='testbot') + bot_name='testbot', + bot_source='bot code location') def test_upload_file(self): client, handler = self._create_client_and_handler_for_file_upload() diff --git a/zulip_bots/zulip_bots/tests/test_run.py b/zulip_bots/zulip_bots/tests/test_run.py index d4096a198..51adf56ac 100644 --- a/zulip_bots/zulip_bots/tests/test_run.py +++ b/zulip_bots/zulip_bots/tests/test_run.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 import os import sys +import entrypoints import zulip_bots.run from zulip_bots.lib import extract_query_without_mention import unittest @@ -15,6 +16,8 @@ class TestDefaultArguments(TestCase): our_dir = os.path.dirname(__file__) path_to_bot = os.path.abspath(os.path.join(our_dir, '../bots/giphy/giphy.py')) + packaged_bot_distro = entrypoints.Distribution("packaged-bot-source", "1.0.0") + packaged_bot_entrypoint = entrypoints.EntryPoint("packaged_bot", "module_name", None, distro=packaged_bot_distro) @patch('sys.argv', ['zulip-run-bot', 'giphy', '--config-file', '/foo/bar/baz.conf']) @patch('zulip_bots.run.run_message_handler_for_bot') @@ -26,6 +29,7 @@ def test_argument_parsing_with_bot_name(self, mock_run_message_handler_for_bot: config_file='/foo/bar/baz.conf', bot_config_file=None, lib_module=mock.ANY, + bot_source='source', quiet=False) @patch('sys.argv', ['zulip-run-bot', path_to_bot, '--config-file', '/foo/bar/baz.conf']) @@ -39,6 +43,23 @@ def test_argument_parsing_with_bot_path(self, mock_run_message_handler_for_bot: config_file='/foo/bar/baz.conf', bot_config_file=None, lib_module=mock.ANY, + bot_source='source', + quiet=False) + + @patch('sys.argv', ['zulip-run-bot', 'packaged_bot', '--config-file', '/foo/bar/baz.conf']) + @patch('zulip_bots.run.run_message_handler_for_bot') + def test_argument_parsing_with_zulip_bot_registry(self, mock_run_message_handler_for_bot: mock.Mock) -> None: + with patch('zulip_bots.run.exit_gracefully_if_zulip_config_is_missing'): + with patch('zulip_bots.finder.entrypoints.EntryPoint.load'): + with patch('zulip_bots.finder.entrypoints.get_group_all', return_value = [self.packaged_bot_entrypoint]): + zulip_bots.run.main() + + mock_run_message_handler_for_bot.assert_called_with( + bot_name='packaged_bot', + config_file='/foo/bar/baz.conf', + bot_config_file=None, + lib_module=mock.ANY, + bot_source='packaged-bot-source: 1.0.0', quiet=False) def test_adding_bot_parent_dir_to_sys_path_when_bot_name_specified(self) -> None: