From 2211acc81d2dd039f576fd0765ec48c79474c6e9 Mon Sep 17 00:00:00 2001 From: Ben Kurtovic Date: Fri, 6 Apr 2012 15:18:15 -0400 Subject: [PATCH] Update TaskManager implementation --- earwigbot/bot.py | 24 +++--- earwigbot/commands/__init__.py | 10 ++- earwigbot/tasks/__init__.py | 167 +++++++++++++++++++++----------------- earwigbot/tasks/afc_catdelink.py | 2 +- earwigbot/tasks/afc_copyvios.py | 2 +- earwigbot/tasks/afc_dailycats.py | 2 +- earwigbot/tasks/afc_history.py | 2 +- earwigbot/tasks/afc_statistics.py | 2 +- earwigbot/tasks/afc_undated.py | 2 +- earwigbot/tasks/blptag.py | 2 +- earwigbot/tasks/feed_dailycats.py | 2 +- earwigbot/tasks/wrongmime.py | 2 +- 12 files changed, 121 insertions(+), 98 deletions(-) diff --git a/earwigbot/bot.py b/earwigbot/bot.py index c6da54b..42f3b89 100644 --- a/earwigbot/bot.py +++ b/earwigbot/bot.py @@ -20,13 +20,13 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -import threading +from threading import Lock, Thread from time import sleep, time from earwigbot.commands import CommandManager from earwigbot.config import BotConfig from earwigbot.irc import Frontend, Watcher -from earwigbot.tasks import task_manager +from earwigbot.tasks import TaskManager __all__ = ["Bot"] @@ -50,29 +50,29 @@ class Bot(object): self.config = BotConfig(root_dir) self.logger = logging.getLogger("earwigbot") self.commands = CommandManager(self) - self.tasks = None + self.tasks = TaskManager(self) self.frontend = None self.watcher = None - self.component_lock = threading.Lock() + self.component_lock = Lock() self._keep_looping = True def _start_irc_components(self): if self.config.components.get("irc_frontend"): self.logger.info("Starting IRC frontend") self.frontend = Frontend(self) - threading.Thread(name=name, target=self.frontend.loop).start() + Thread(name=name, target=self.frontend.loop).start() if self.config.components.get("irc_watcher"): self.logger.info("Starting IRC watcher") self.watcher = Watcher(self) - threading.Thread(name=name, target=self.watcher.loop).start() + Thread(name=name, target=self.watcher.loop).start() def _start_wiki_scheduler(self): def wiki_scheduler(): while self._keep_looping: time_start = time() - task_manager.schedule() + self.tasks.schedule() time_end = time() time_diff = time_start - time_end if time_diff < 60: # Sleep until the next minute @@ -80,7 +80,7 @@ class Bot(object): if self.config.components.get("wiki_scheduler"): self.logger.info("Starting wiki scheduler") - threading.Thread(name=name, target=wiki_scheduler).start() + Thread(name=name, target=wiki_scheduler).start() def _stop_irc_components(self): if self.frontend: @@ -93,10 +93,10 @@ class Bot(object): with self.component_lock: if self.frontend and self.frontend.is_stopped(): self.frontend = Frontend(self) - threading.Thread(name=name, target=self.frontend.loop).start() + Thread(name=name, target=self.frontend.loop).start() if self.watcher and self.watcher.is_stopped(): self.watcher = Watcher(self) - threading.Thread(name=name, target=self.watcher.loop).start() + Thread(name=name, target=self.watcher.loop).start() sleep(5) def run(self): @@ -106,7 +106,8 @@ class Bot(object): self.config.decrypt(config.wiki, "search", "credentials", "secret") self.config.decrypt(config.irc, "frontend", "nickservPassword") self.config.decrypt(config.irc, "watcher", "nickservPassword") - self.commands.load() + self.commands.load() + self.tasks.load() self._start_irc_components() self._start_wiki_scheduler() self._loop() @@ -116,6 +117,7 @@ class Bot(object): self._stop_irc_components() self.config.load() self.commands.load() + self.tasks.load() self._start_irc_components() def stop(self): diff --git a/earwigbot/commands/__init__.py b/earwigbot/commands/__init__.py index e97d2b9..a22fcf6 100644 --- a/earwigbot/commands/__init__.py +++ b/earwigbot/commands/__init__.py @@ -115,9 +115,15 @@ class CommandManager(object): f.close() try: - command = module.Command(self.bot) + command_class = module.Command except AttributeError: return # No command in this module + try: + command = command_class(self.bot) + except Exception: + e = "Error initializing Command() class in {0} (from {1})" + self.logger.exception(e.format(name, path)) + return if not isinstance(command, BaseCommand): return @@ -129,7 +135,7 @@ class CommandManager(object): with self._command_access_lock: self._commands.clear() dirs = [path.join(path.dirname(__file__), "commands"), - path.join(bot.config.root_dir, "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] != "_"] diff --git a/earwigbot/tasks/__init__.py b/earwigbot/tasks/__init__.py index 7e57b30..dcd963c 100644 --- a/earwigbot/tasks/__init__.py +++ b/earwigbot/tasks/__init__.py @@ -25,39 +25,40 @@ EarwigBot's Wiki Task Manager This package provides the wiki bot "tasks" EarwigBot runs. This module contains the BaseTask class (import with `from earwigbot.tasks import BaseTask`) and an -internal _TaskManager class. This can be accessed through the `task_manager` -singleton. +internal TaskManager class. This can be accessed through `bot.tasks`. """ -import logging -import os -import sys -import threading -import time +import imp +from os import listdir, path +from threading import Lock, Thread +from time import gmtime, strftime from earwigbot import wiki -from earwigbot.config import config -__all__ = ["BaseTask", "task_manager"] +__all__ = ["BaseTask", "TaskManager"] class BaseTask(object): """A base class for bot tasks that edit Wikipedia.""" name = None number = 0 - def __init__(self): + def __init__(self, bot): """Constructor for new tasks. This is called once immediately after the task class is loaded by - the task manager (in tasks._load_task()). + the task manager (in tasks._load_task()). Don't override this directly + (or if you do, remember super(Task, self).__init()) - use setup(). """ - pass + self.bot = bot + self.logger = bot.tasks.logger.getLogger(self.name) + self.setup() + + def setup(self): + """Hook called immediately after the task is loaded. - def _setup_logger(self): - """Set up a basic module-level logger.""" - logger_name = ".".join(("earwigbot", "tasks", self.name)) - self.logger = logging.getLogger(logger_name) - self.logger.setLevel(logging.DEBUG) + Does nothing by default; feel free to override. + """ + pass def run(self, **kwargs): """Main entry point to run a given task. @@ -83,7 +84,7 @@ class BaseTask(object): If the config value is not found, we just return the arg as-is. """ try: - summary = config.wiki["summary"] + summary = self.bot.config.wiki["summary"] except KeyError: return comment return summary.replace("$1", str(self.number)).replace("$2", comment) @@ -111,7 +112,7 @@ class BaseTask(object): site = wiki.get_site() try: - cfg = config.wiki["shutoff"] + cfg = self.bot.config.wiki["shutoff"] except KeyError: return False title = cfg.get("page", "User:$1/Shutoff/Task $2") @@ -130,91 +131,107 @@ class BaseTask(object): return True -class _TaskManager(object): - def __init__(self): - self.logger = logging.getLogger("earwigbot.commands") - self._base_dir = os.path.dirname(os.path.abspath(__file__)) +class TaskManager(object): + def __init__(self, bot): + self.bot = bot + self.logger = bot.logger.getLogger("tasks") self._tasks = {} + self._task_access_lock = Lock() + + def _wrapper(self, task, **kwargs): + """Wrapper for task classes: run the task and catch any errors.""" + try: + task.run(**kwargs) + except Exception: + msg = "Task '{0}' raised an exception and had to stop:" + self.logger.exception(msg.format(task.name)) + else: + msg = "Task '{0}' finished without error" + self.logger.info(msg.format(task.name)) + + def _load_task(self, name, path): + """Load a specific task from a module, identified by name and path. - def _load_task(self, filename): - """Load a specific task from a module, identified by file name.""" - # Strip .py from the filename's end and join with our package name: - name = ".".join(("tasks", filename[:-3])) + We'll first try to import it using imp magic, and if that works, make + an instance of the 'Task' class inside (assuming it is an instance of + BaseTask), add it to self._tasks, and log the addition. Any problems + along the way will either be ignored or logged. + """ + f, path, desc = imp.find_module(name, [path]) try: - __import__(name) - except: - self.logger.exception("Couldn't load file {0}:".format(filename)) + module = imp.load_module(name, f, path, desc) + except Exception: + e = "Couldn't load module {0} from {1}" + self.logger.exception(e.format(name, path)) return + finally: + f.close() try: - task = sys.modules[name].Task() + task_class = module.Task except AttributeError: return # No task in this module + try: + task = task_class(self.bot) + except Exception: + e = "Error initializing Task() class in {0} (from {1})" + self.logger.exception(e.format(name, path)) + return if not isinstance(task, BaseTask): return - task._setup_logger() self._tasks[task.name] = task self.logger.debug("Added task {0}".format(task.name)) - def _wrapper(self, task, **kwargs): - """Wrapper for task classes: run the task and catch any errors.""" - try: - task.run(**kwargs) - except: - msg = "Task '{0}' raised an exception and had to stop" - self.logger.exception(msg.format(task.name)) - else: - msg = "Task '{0}' finished without error" - self.logger.info(msg.format(task.name)) - def load(self): - """Load all valid tasks from tasks/ into self._tasks.""" - files = os.listdir(self._base_dir) - files.sort() - - for filename in files: - if filename.startswith("_") or not filename.endswith(".py"): - continue - self._load_task(filename) + """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}" tasks = ', '.join(self._tasks.keys()) self.logger.info(msg.format(len(self._tasks), tasks)) + def start(self, task_name, **kwargs): + """Start a given task in a new thread. kwargs are passed to task.run""" + msg = "Starting task '{0}' in a new thread" + self.logger.info(msg.format(task_name)) + + with self._task_access_lock: + try: + task = self._tasks[task_name] + except KeyError: + e = "Couldn't find task '{0}':" + self.logger.error(e.format(task_name)) + return + + task_thread = Thread(target=self._wrapper, args=(task,), kwargs=kwargs) + start_time = strftime("%b %d %H:%M:%S") + task_thread.name = "{0} ({1})".format(task_name, start_time) + task_thread.start() + def schedule(self, now=None): """Start all tasks that are supposed to be run at a given time.""" if not now: - now = time.gmtime() + now = gmtime() # Get list of tasks to run this turn: - tasks = config.schedule(now.tm_min, now.tm_hour, now.tm_mday, - now.tm_mon, now.tm_wday) + tasks = self.bot.config.schedule(now.tm_min, now.tm_hour, now.tm_mday, + now.tm_mon, now.tm_wday) for task in tasks: if isinstance(task, list): # They've specified kwargs, - self.start(task[0], **task[1]) # so pass those to start_task + self.start(task[0], **task[1]) # so pass those to start else: # Otherwise, just pass task_name self.start(task) - def start(self, task_name, **kwargs): - """Start a given task in a new thread. Pass args to the task's run() - function.""" - msg = "Starting task '{0}' in a new thread" - self.logger.info(msg.format(task_name)) - - try: - task = self._tasks[task_name] - except KeyError: - e = "Couldn't find task '{0}': bot/tasks/{0}.py does not exist" - self.logger.error(e.format(task_name)) - return - - func = lambda: self._wrapper(task, **kwargs) - task_thread = threading.Thread(target=func) - start_time = time.strftime("%b %d %H:%M:%S") - task_thread.name = "{0} ({1})".format(task_name, start_time) - task_thread.start() - def get(self, task_name): """Return the class instance associated with a certain task name. @@ -225,5 +242,3 @@ class _TaskManager(object): def get_all(self): """Return our dict of all loaded tasks.""" return self._tasks - -task_manager = _TaskManager() diff --git a/earwigbot/tasks/afc_catdelink.py b/earwigbot/tasks/afc_catdelink.py index c5d3c0f..7888e70 100644 --- a/earwigbot/tasks/afc_catdelink.py +++ b/earwigbot/tasks/afc_catdelink.py @@ -27,7 +27,7 @@ class Task(BaseTask): submissions.""" name = "afc_catdelink" - def __init__(self): + def setup(self): pass def run(self, **kwargs): diff --git a/earwigbot/tasks/afc_copyvios.py b/earwigbot/tasks/afc_copyvios.py index 4db5a63..0956b43 100644 --- a/earwigbot/tasks/afc_copyvios.py +++ b/earwigbot/tasks/afc_copyvios.py @@ -36,7 +36,7 @@ class Task(BaseTask): name = "afc_copyvios" number = 1 - def __init__(self): + def setup(self): cfg = config.tasks.get(self.name, {}) self.template = cfg.get("template", "AfC suspected copyvio") self.ignore_list = cfg.get("ignoreList", []) diff --git a/earwigbot/tasks/afc_dailycats.py b/earwigbot/tasks/afc_dailycats.py index efddd20..b286924 100644 --- a/earwigbot/tasks/afc_dailycats.py +++ b/earwigbot/tasks/afc_dailycats.py @@ -27,7 +27,7 @@ class Task(BaseTask): name = "afc_dailycats" number = 3 - def __init__(self): + def setup(self): pass def run(self, **kwargs): diff --git a/earwigbot/tasks/afc_history.py b/earwigbot/tasks/afc_history.py index 03117ad..9e146f9 100644 --- a/earwigbot/tasks/afc_history.py +++ b/earwigbot/tasks/afc_history.py @@ -57,7 +57,7 @@ class Task(BaseTask): """ name = "afc_history" - def __init__(self): + def setup(self): cfg = config.tasks.get(self.name, {}) self.num_days = cfg.get("days", 90) self.categories = cfg.get("categories", {}) diff --git a/earwigbot/tasks/afc_statistics.py b/earwigbot/tasks/afc_statistics.py index 3de023d..79070ae 100644 --- a/earwigbot/tasks/afc_statistics.py +++ b/earwigbot/tasks/afc_statistics.py @@ -53,7 +53,7 @@ class Task(BaseTask): name = "afc_statistics" number = 2 - def __init__(self): + def setup(self): self.cfg = cfg = config.tasks.get(self.name, {}) # Set some wiki-related attributes: diff --git a/earwigbot/tasks/afc_undated.py b/earwigbot/tasks/afc_undated.py index 512f09d..e596d1d 100644 --- a/earwigbot/tasks/afc_undated.py +++ b/earwigbot/tasks/afc_undated.py @@ -26,7 +26,7 @@ class Task(BaseTask): """A task to clear [[Category:Undated AfC submissions]].""" name = "afc_undated" - def __init__(self): + def setup(self): pass def run(self, **kwargs): diff --git a/earwigbot/tasks/blptag.py b/earwigbot/tasks/blptag.py index 5bb4052..505b7d9 100644 --- a/earwigbot/tasks/blptag.py +++ b/earwigbot/tasks/blptag.py @@ -27,7 +27,7 @@ class Task(BaseTask): {{WP Biography}}.""" name = "blptag" - def __init__(self): + def setup(self): pass def run(self, **kwargs): diff --git a/earwigbot/tasks/feed_dailycats.py b/earwigbot/tasks/feed_dailycats.py index 3d6afd7..a068af6 100644 --- a/earwigbot/tasks/feed_dailycats.py +++ b/earwigbot/tasks/feed_dailycats.py @@ -26,7 +26,7 @@ class Task(BaseTask): """A task to create daily categories for [[WP:FEED]].""" name = "feed_dailycats" - def __init__(self): + def setup(self): pass def run(self, **kwargs): diff --git a/earwigbot/tasks/wrongmime.py b/earwigbot/tasks/wrongmime.py index 1b51589..ed531d5 100644 --- a/earwigbot/tasks/wrongmime.py +++ b/earwigbot/tasks/wrongmime.py @@ -27,7 +27,7 @@ class Task(BaseTask): type.""" name = "wrongmime" - def __init__(self): + def setup(self): pass def run(self, **kwargs):