Some polishing work on celeryctl

celery/bin/celeryctl.py

@@ -1,3 +1,4 @@
+import os
 import sys
 
 from optparse import OptionParser, make_option as Option
@@ -8,6 +9,14 @@ from anyjson import deserialize
 from celery import __version__
 
 
+commands = {}
+
+
+def command(fun):
+    commands[fun.__name__] = fun
+    return fun
+
+
 class Command(object):
     help = ""
     args = ""
@@ -18,15 +27,18 @@ class Command(object):
             help="Celery loaders module name (default: default)"),
     )
 
-    def create_parser(self, prog_name, subcommand):
+    def create_parser(self, prog_name, command):
         return OptionParser(prog=prog_name,
-                            usage=self.usage(subcommand),
+                            usage=self.usage(command),
                             version=self.get_version(),
                             option_list=self.option_list)
 
     def run_from_argv(self, argv):
-        parser = self.create_parser(argv[0], argv[1])
-        options, args = parser.parse_args(argv[2:])
+        self.prog_name = os.path.basename(argv[0])
+        self.command = argv[1]
+        self.arglist = argv[2:]
+        self.parser = self.create_parser(self.prog_name, self.command)
+        options, args = self.parser.parse_args(self.arglist)
         if options.loader:
             os.environ["CELERY_LOADER"] = options.loader
         if options.conf:
@@ -44,11 +56,12 @@ class Command(object):
     def get_version(self):
         return __version__
 
-    def usage(self, subcommand):
-        return "%%prog %s [options] %s" % (subcommand, self.args)
+    def usage(self, command):
+        return "%%prog %s [options] %s" % (command, self.args)
 
 
 class apply(Command):
+    args = "<task_name>"
     option_list = Command.option_list + (
             Option("--args", "-a", dest="args"),
             Option("--kwargs", "-k", dest="kwargs"),
@@ -91,6 +104,7 @@ class apply(Command):
                         expires=expires)
 
         print(res.task_id)
+apply = command(apply)
 
 class inspect(Command):
     choices = ("active", "scheduled", "reserved",
@@ -103,9 +117,9 @@ class inspect(Command):
                 Option("--destination", "-d", dest="destination",
                     help="Comma separated list of destination node names."))
 
-    def usage(self, subcommand):
+    def usage(self, command):
         return "%%prog %s [options] %s [%s]" % (
-                subcommand, self.args, "|".join(self.choices))
+                command, self.args, "|".join(self.choices))
 
     def run(self, *args, **kwargs):
         if not args:
@@ -122,15 +136,57 @@ class inspect(Command):
 
         i = inspect(destination=destination, timeout=timeout)
         pprint(getattr(i, command)())
+inspect = command(inspect)
+
+
+class help(Command):
 
+    def usage(self, command):
+        return "%%prog <command> [options] %s" % (self.args, )
 
-commands = {"apply": apply,
-            "inspect": inspect}
+    def run(self, *args, **kwargs):
+        self.parser.print_help()
+        usage = ["",
+                "Type '%s <command> --help' for help on a "
+                "specific command." % (self.prog_name, ),
+                "",
+                "Available commands:"]
+        for command in list(sorted(commands.keys())):
+            usage.append("    %s" % command)
+        print("\n".join(usage))
+help = command(help)
+
+
+class celeryctl(object):
+    commands = commands
+
+    def execute(self, command, argv=None):
+        if argv is None:
+            argv = sys.arg
+        argv = list(argv)
+        try:
+            cls = self.commands[command]
+        except KeyError:
+            cls = self.commands["help"]
+            argv.insert(1, "help")
+        cls = self.commands.get(command) or self.commands["help"]
+        try:
+            cls().run_from_argv(argv)
+        except TypeError:
+            return self.execute("help", argv)
+
+    def execute_from_commandline(self, argv=None):
+        if argv is None:
+            argv = sys.argv
+        try:
+            command = argv[1]
+        except IndexError:
+            command = "help"
+        return self.execute(command, argv)
 
 
 def main():
-    subcommand = sys.argv[1]
-    commands[subcommand]().run_from_argv(sys.argv)
+    celeryctl().execute_from_commandline()
 
 if __name__ == "__main__":
     main()