Update TaskManager implementation

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

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] != "_"]

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

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

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", [])

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

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", {})

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:

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

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

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

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