Bots: Add support for running packaged bots via entrypoints

zulip_bots/setup.py

@@ -55,6 +55,7 @@
         'html2text',
         'lxml',
         'BeautifulSoup4',
+        'entrypoints',
     ],
 )
 

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)

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())

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('''

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()

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: