From 60b52f6aeb623fcff077e93181c983d47ada74e1 Mon Sep 17 00:00:00 2001
From: Ben Kurtovic <ben.kurtovic@verizon.net>
Date: Fri, 6 Apr 2012 16:12:28 -0400
Subject: [PATCH] __iter__ for Command+TaskManager instead of get_all()

---
 earwigbot/commands/__init__.py | 17 +++++++++++++----
 earwigbot/commands/help.py     |  8 ++++----
 earwigbot/commands/threads.py  | 10 +++-------
 earwigbot/tasks/__init__.py    |  8 ++++----
 4 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/earwigbot/commands/__init__.py b/earwigbot/commands/__init__.py
index ba09eac..4531abc 100644
--- a/earwigbot/commands/__init__.py
+++ b/earwigbot/commands/__init__.py
@@ -71,6 +71,8 @@ class BaseCommand(object):
 
         Given a Data() instance, return True if we should respond to this
         activity, or False if we should ignore it or it doesn't apply to us.
+        Be aware that since this is called for each message sent on IRC, it
+        should not be cheap to execute and unlikely to throw exceptions.
 
         Most commands return True if data.command == self.name, otherwise they
         return False. This is the default behavior of check(); you need only
@@ -97,6 +99,10 @@ class CommandManager(object):
         self._commands = {}
         self._command_access_lock = Lock()
 
+    def __iter__(self):
+        for name in self._commands:
+            yield name
+
     def _load_command(self, name, path):
         """Load a specific command from a module, identified by name and path.
 
@@ -148,10 +154,6 @@ class CommandManager(object):
         commands = ", ".join(self._commands.keys())
         self.logger.info(msg.format(len(self._commands), commands))
 
-    def get_all(self):
-        """Return our dict of all loaded commands."""
-        return self._commands
-
     def check(self, hook, data):
         """Given an IRC event, check if there's anything we can respond to."""
         with self._command_access_lock:
@@ -164,3 +166,10 @@ class CommandManager(object):
                             e = "Error executing command '{0}':"
                             self.logger.exception(e.format(data.command))
                         break
+
+    def get(self, command_name):
+        """Return the class instance associated with a certain command name.
+
+        Will raise KeyError if the command is not found.
+        """
+        return self._command[command_name]
diff --git a/earwigbot/commands/help.py b/earwigbot/commands/help.py
index bca6ee2..95d9e39 100644
--- a/earwigbot/commands/help.py
+++ b/earwigbot/commands/help.py
@@ -30,7 +30,6 @@ class Command(BaseCommand):
     name = "help"
 
     def process(self, data):
-        self.cmnds = self.bot.commands.get_all()
         if not data.args:
             self.do_main_help(data)
         else:
@@ -39,7 +38,7 @@ class Command(BaseCommand):
     def do_main_help(self, data):
         """Give the user a general help message with a list of all commands."""
         msg = "Hi, I'm a bot! I have {0} commands loaded: {1}. You can get help for any command with '!help <command>'."
-        cmnds = sorted(self.cmnds.keys())
+        cmnds = sorted(self.bot.commands)
         msg = msg.format(len(cmnds), ', '.join(cmnds))
         self.connection.reply(data, msg)
 
@@ -53,13 +52,14 @@ class Command(BaseCommand):
         dummy.command = command.lower()
         dummy.is_command = True
 
-        for cmnd in self.cmnds.values():
+        for cmnd_name in self.bot.commands:
+            cmnd = self.bot.commands.get(cmnd_name)
             if not cmnd.check(dummy):
                 continue
             if cmnd.__doc__:
                 doc = cmnd.__doc__.replace("\n", "")
                 doc = re.sub("\s\s+", " ", doc)
-                msg = "info for command \x0303{0}\x0301: \"{1}\""
+                msg = "help for command \x0303{0}\x0301: \"{1}\""
                 self.connection.reply(data, msg.format(command, doc))
                 return
             break
diff --git a/earwigbot/commands/threads.py b/earwigbot/commands/threads.py
index 13d0a70..bfe894b 100644
--- a/earwigbot/commands/threads.py
+++ b/earwigbot/commands/threads.py
@@ -103,13 +103,9 @@ class Command(BaseCommand):
     def do_listall(self):
         """With !tasks listall or !tasks all, list all loaded tasks, and report
         whether they are currently running or idle."""
-        all_tasks = self.bot.tasks.get_all().keys()
         threads = threading.enumerate()
         tasklist = []
-
-        all_tasks.sort()
-
-        for task in all_tasks:
+        for task in sorted(self.bot.tasks):
             threadlist = [t for t in threads if t.name.startswith(task)]
             ids = [str(t.ident) for t in threadlist]
             if not ids:
@@ -144,9 +140,9 @@ class Command(BaseCommand):
             self.connection.reply(data, msg)
             return
 
-        if task_name not in self.bot.tasks.get_all().keys():
+        if task_name not in self.bot.tasks:
             # This task does not exist or hasn't been loaded:
-            msg = "task could not be found; either tasks/{0}.py doesn't exist, or it wasn't loaded correctly."
+            msg = "task could not be found; either it doesn't exist, or it wasn't loaded correctly."
             self.connection.reply(data, msg.format(task_name))
             return
 
diff --git a/earwigbot/tasks/__init__.py b/earwigbot/tasks/__init__.py
index 7c268f9..a1f75f3 100644
--- a/earwigbot/tasks/__init__.py
+++ b/earwigbot/tasks/__init__.py
@@ -139,6 +139,10 @@ class TaskManager(object):
         self._tasks = {}
         self._task_access_lock = Lock()
 
+    def __iter__(self):
+        for name in self._tasks:
+            yield name
+
     def _wrapper(self, task, **kwargs):
         """Wrapper for task classes: run the task and catch any errors."""
         try:
@@ -239,7 +243,3 @@ class TaskManager(object):
         Will raise KeyError if the task is not found.
         """
         return self._tasks[task_name]
-
-    def get_all(self):
-        """Return our dict of all loaded tasks."""
-        return self._tasks