From 03062e808b67e5e16c4c0c33e197468ff5678d3b Mon Sep 17 00:00:00 2001 From: Ben Kurtovic Date: Sat, 7 Apr 2012 17:41:33 -0400 Subject: [PATCH] Wrote the command-line utility, added logging levels, improved Bot organization - Fixed loading bugs in CommandLoader and TaskLoader --- earwigbot/bot.py | 21 +++++-------- earwigbot/commands/__init__.py | 31 ++++++++++++------- earwigbot/config.py | 31 ++++++++++++++----- earwigbot/tasks/__init__.py | 32 ++++++++++++------- earwigbot/util.py | 70 +++++++++++++++++++++++------------------- 5 files changed, 112 insertions(+), 73 deletions(-) diff --git a/earwigbot/bot.py b/earwigbot/bot.py index ea40dcd..8bf92bb 100644 --- a/earwigbot/bot.py +++ b/earwigbot/bot.py @@ -53,8 +53,8 @@ class Bot(object): wiki toolset with bot.wiki.get_site(). """ - def __init__(self, root_dir): - self.config = BotConfig(root_dir) + def __init__(self, root_dir, level=logging.INFO): + self.config = BotConfig(root_dir, level) self.logger = logging.getLogger("earwigbot") self.commands = CommandManager(self) self.tasks = TaskManager(self) @@ -65,6 +65,10 @@ class Bot(object): self.component_lock = Lock() self._keep_looping = True + self.config.load() + self.commands.load() + self.tasks.load() + def _start_irc_components(self): if self.config.components.get("irc_frontend"): self.logger.info("Starting IRC frontend") @@ -107,18 +111,10 @@ class Bot(object): self.logger.warn("IRC watcher has stopped; restarting") self.watcher = Watcher(self) Thread(name=name, target=self.watcher.loop).start() - sleep(5) + sleep(3) def run(self): - config = self.config - config.load() - config.decrypt(config.wiki, "password") - config.decrypt(config.wiki, "search", "credentials", "key") - config.decrypt(config.wiki, "search", "credentials", "secret") - config.decrypt(config.irc, "frontend", "nickservPassword") - config.decrypt(config.irc, "watcher", "nickservPassword") - self.commands.load() - self.tasks.load() + self.logger.info("Starting bot") self._start_irc_components() self._start_wiki_scheduler() self._loop() @@ -137,4 +133,3 @@ class Bot(object): with self.component_lock: self._stop_irc_components() self._keep_looping = False - sleep(3) # Give a few seconds to finish closing IRC connections diff --git a/earwigbot/commands/__init__.py b/earwigbot/commands/__init__.py index dfe2312..98cebec 100644 --- a/earwigbot/commands/__init__.py +++ b/earwigbot/commands/__init__.py @@ -135,22 +135,31 @@ class CommandManager(object): return self._commands[command.name] = command - self.logger.debug("Added command {0}".format(command.name)) + self.logger.debug("Loaded command {0}".format(command.name)) + + def _load_directory(self, dir): + """Load all valid commands in a given directory.""" + processed = [] + for name in listdir(dir): + if not name.endswith(".py") and not name.endswith(".pyc"): + continue + if name.startswith("_") or name.startswith("."): + continue + modname = sub("\.pyc?$", "", name) # Remove extension + if modname not in processed: + self._load_command(modname, dir) + processed.append(modname) def load(self): """Load (or reload) all valid commands into self._commands.""" with self._command_access_lock: self._commands.clear() - dirs = [path.join(path.dirname(__file__), "commands"), - path.join(self.bot.config.root_dir, "commands")] - for dir in dirs: - files = listdir(dir) - files = [sub("\.pyc?$", "", f) for f in files if f[0] != "_"] - files = list(set(files)) # Remove duplicates - for filename in sorted(files): - self._load_command(filename, dir) - - msg = "Found {0} commands: {1}" + builtin_dir = path.dirname(__file__) + plugins_dir = path.join(self.bot.config.root_dir, "commands") + self._load_directory(builtin_dir) # Built-in commands + self._load_directory(plugins_dir) # Custom commands, aka plugins + + msg = "Loaded {0} commands: {1}" commands = ", ".join(self._commands.keys()) self.logger.info(msg.format(len(self._commands), commands)) diff --git a/earwigbot/config.py b/earwigbot/config.py index fc73a08..962f8d3 100644 --- a/earwigbot/config.py +++ b/earwigbot/config.py @@ -59,8 +59,9 @@ class BotConfig(object): aren't encrypted """ - def __init__(self, root_dir): + def __init__(self, root_dir, level): self._root_dir = root_dir + self._logging_level = level self._config_path = path.join(self._root_dir, "config.yml") self._log_dir = path.join(self._root_dir, "logs") self._decryption_key = None @@ -74,7 +75,14 @@ class BotConfig(object): self._nodes = [self._components, self._wiki, self._tasks, self._irc, self._metadata] - self._decryptable_nodes = [] + + self._decryptable_nodes = [ # Default nodes to decrypt + (self._wiki, ("password")), + (self._wiki, ("search", "credentials", "key")), + (self._wiki, ("search", "credentials", "secret")), + (self._irc, ("frontend", "nickservPassword")), + (self._irc, ("watcher", "nickservPassword")), + ] def _load(self): """Load data from our JSON config file (config.yml) into self._data.""" @@ -119,10 +127,10 @@ class BotConfig(object): h.setFormatter(formatter) logger.addHandler(h) - stream_handler = logging.StreamHandler() - stream_handler.setLevel(logging.DEBUG) - stream_handler.setFormatter(color_formatter) - logger.addHandler(stream_handler) + self._stream_handler = stream = logging.StreamHandler() + stream.setLevel(self._logging_level) + stream.setFormatter(color_formatter) + logger.addHandler(stream) def _decrypt(self, node, nodes): """Try to decrypt the contents of a config node. Use self.decrypt().""" @@ -148,6 +156,15 @@ class BotConfig(object): return self._root_dir @property + def logging_level(self): + return self._logging_level + + @logging_level.setter + def logging_level(self, level): + self._logging_level = level + self._stream_handler.setLevel(level) + + @property def path(self): return self._config_path @@ -213,7 +230,7 @@ class BotConfig(object): if choice.lower().startswith("y"): self._make_new() else: - exit(1) + exit(1) # TODO: raise an exception instead self._load() data = self._data diff --git a/earwigbot/tasks/__init__.py b/earwigbot/tasks/__init__.py index aa6e35e..d70bcde 100644 --- a/earwigbot/tasks/__init__.py +++ b/earwigbot/tasks/__init__.py @@ -30,6 +30,7 @@ internal TaskManager class. This can be accessed through `bot.tasks`. import imp from os import listdir, path +from re import sub from threading import Lock, Thread from time import gmtime, strftime @@ -186,22 +187,31 @@ class TaskManager(object): return self._tasks[task.name] = task - self.logger.debug("Added task {0}".format(task.name)) + self.logger.debug("Loaded task {0}".format(task.name)) + + def _load_directory(self, dir): + """Load all valid tasks in a given directory.""" + processed = [] + for name in listdir(dir): + if not name.endswith(".py") and not name.endswith(".pyc"): + continue + if name.startswith("_") or name.startswith("."): + continue + modname = sub("\.pyc?$", "", name) # Remove extension + if modname not in processed: + self._load_task(modname, dir) + processed.append(modname) def load(self): """Load (or reload) all valid tasks into self._tasks.""" with self._task_access_lock: self._tasks.clear() - dirs = [path.join(path.dirname(__file__), "tasks"), - path.join(self.bot.config.root_dir, "tasks")] - for dir in dirs: - files = listdir(dir) - files = [sub("\.pyc?$", "", f) for f in files if f[0] != "_"] - files = list(set(files)) # Remove duplicates - for filename in sorted(files): - self._load_task(filename) - - msg = "Found {0} tasks: {1}" + builtin_dir = path.dirname(__file__) + plugins_dir = path.join(self.bot.config.root_dir, "tasks") + self._load_directory(builtin_dir) # Built-in tasks + self._load_directory(plugins_dir) # Custom tasks, aka plugins + + msg = "Loaded {0} tasks: {1}" tasks = ', '.join(self._tasks.keys()) self.logger.info(msg.format(len(self._tasks), tasks)) diff --git a/earwigbot/util.py b/earwigbot/util.py index 5f139da..442bc6e 100755 --- a/earwigbot/util.py +++ b/earwigbot/util.py @@ -21,46 +21,54 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. +""" +This is EarwigBot's command-line utility, enabling you to easily start the +bot or run specific tasks. +""" + import argparse +import logging from os import path from earwigbot import __version__ from earwigbot.bot import Bot -__all__ = ["BotUtility", "main"] - -class BotUtility(object): - """ - This is a command-line utility for EarwigBot that enables you to easily - start the bot without writing generally unnecessary three-line bootstrap - scripts. It supports starting the bot from any directory, as well as - starting individual tasks instead of the entire bot. - """ - - def version(self): - return "EarwigBot v{0}".format(__version__) - - def run(self, root_dir): - bot = Bot(root_dir) - print self.version() - #try: - # bot.run() - #finally: - # bot.stop() - - def main(self): - parser = argparse.ArgumentParser(description=BotUtility.__doc__) - - parser.add_argument("-v", "--version", action="version", - version=self.version()) - parser.add_argument("root_dir", metavar="path", nargs="?", default=path.curdir) - args = parser.parse_args() +__all__ = ["main"] - root_dir = path.abspath(args.root_dir) - self.run(root_dir) +def main(): + version = "EarwigBot v{0}".format(__version__) + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("path", nargs="?", metavar="PATH", default=path.curdir, + help="path to the bot's working directory, which will be created if it doesn't exist; current directory assumed if not specified") + parser.add_argument("-v", "--version", action="version", version=version) + parser.add_argument("-d", "--debug", action="store_true", + help="print all logs, including DEBUG-level messages") + parser.add_argument("-q", "--quiet", action="store_true", + help="don't print any logs except warnings and errors") + parser.add_argument("-t", "--task", metavar="NAME", + help="given the name of a task, the bot will run it instead of the main bot and then exit") + args = parser.parse_args() + if args.debug and args.quiet: + parser.print_usage() + print "earwigbot: error: cannot show debug messages and be quiet at the same time" + return + level = logging.INFO + if args.debug: + level = logging.DEBUG + elif args.quiet: + level = logging.WARNING -main = BotUtility().main + print version + print + bot = Bot(path.abspath(args.path), level=level) + try: + if args.task: + bot.tasks.start(args.task) + else: + bot.run() + finally: + bot.stop() if __name__ == "__main__": main()