Rewrite wagtail command to use argparse for each subcommand

wagtail/bin/wagtail.py

@@ -2,42 +2,59 @@
 import fileinput
 import os
 import re
-from optparse import OptionParser
+import sys
+from argparse import ArgumentParser
 
 from django.core.management import ManagementUtility
 
 
 class Command:
     description = None
-    usage = None
+
+    def create_parser(self, command_name=None):
+        if command_name is None:
+            prog = None
+        else:
+            # hack the prog name as reported to ArgumentParser to include the command
+            prog = "%s %s" % (prog_name(), command_name)
+
+        parser = ArgumentParser(
+            description=getattr(self, 'description', None), add_help=False, prog=prog
+        )
+        self.add_arguments(parser)
+        return parser
+
+    def add_arguments(self, parser):
+        pass
+
+    def print_help(self, command_name):
+        parser = self.create_parser(command_name=command_name)
+        parser.print_help()
+
+    def execute(self, argv):
+        parser = self.create_parser()
+        options = parser.parse_args(sys.argv[2:])
+        options_dict = vars(options)
+        self.run(**options_dict)
 
 
 class CreateProject(Command):
     description = "Creates the directory structure for a new Wagtail project."
-    usage = "Usage: %prog start project_name [directory]"
 
-    def run(self, parser, options, args):
-        # Validate args
-        if len(args) < 2:
-            parser.error("Please specify a name for your Wagtail installation")
-        elif len(args) > 3:
-            parser.error("Too many arguments")
-
-        project_name = args[1]
-        try:
-            dest_dir = args[2]
-        except IndexError:
-            dest_dir = None
+    def add_arguments(self, parser):
+        parser.add_argument('project_name', help="Name for your Wagtail project")
+        parser.add_argument('dest_dir', nargs='?', help="Destination directory inside which to create the project")
 
+    def run(self, project_name=None, dest_dir=None):
         # Make sure given name is not already in use by another python package/module.
         try:
             __import__(project_name)
         except ImportError:
             pass
         else:
-            parser.error("'%s' conflicts with the name of an existing "
-                         "Python module and cannot be used as a project "
-                         "name. Please try another name." % project_name)
+            sys.exit("'%s' conflicts with the name of an existing "
+                     "Python module and cannot be used as a project "
+                     "name. Please try another name." % project_name)
 
         print("Creating a Wagtail project called %(project_name)s" % {'project_name': project_name})  # noqa
 
@@ -66,7 +83,6 @@ class CreateProject(Command):
 
 class UpdateModulePaths(Command):
     description = "Update a Wagtail project tree to use Wagtail 2.x module paths"
-    usage = "Usage: %prog updatemodulepaths [root-path]"
 
     REPLACEMENTS = [
         (re.compile(r'\bwagtail\.wagtailcore\b'), 'wagtail.core'),
@@ -87,14 +103,11 @@ class UpdateModulePaths(Command):
         (re.compile(r'\bwagtail\.contrib\.wagtailstyleguide\b'), 'wagtail.contrib.styleguide'),
     ]
 
-    def run(self, parser, options, args):
-        # Validate args
-        if len(args) > 2:
-            parser.error("Too many arguments")
+    def add_arguments(self, parser):
+        parser.add_argument('root_path', nargs='?', help="Path to your project's root")
 
-        try:
-            root_path = args[1]
-        except IndexError:
+    def run(self, root_path=None):
+        if root_path is None:
             root_path = os.getcwd()
 
         checked_count = 0
@@ -135,36 +148,53 @@ COMMANDS = {
 }
 
 
+def prog_name():
+    return os.path.basename(sys.argv[0])
+
+
+def help_index():
+    print("Type '%s help <subcommand>' for help on a specific subcommand.\n" % prog_name())  # NOQA
+    print("Available subcommands:\n")  # NOQA
+    for name, cmd in sorted(COMMANDS.items()):
+        print("    %s%s" % (name.ljust(20), cmd.description))  # NOQA
+
+
+def unknown_command(command):
+    print("Unknown command: '%s'" % command)  # NOQA
+    print("Type '%s help' for usage." % prog_name())  # NOQA
+    sys.exit(1)
+
+
 def main():
-    # Set up usage string
-    command_descriptions = '\n'.join([
-        "    %s%s" % (name.ljust(20), cmd.description)
-        for name, cmd in sorted(COMMANDS.items())
-        if cmd.description is not None
-    ])
-    usage = (
-        "Usage: %prog <command> [command-options]\n\nAvailable commands:\n\n" +
-        command_descriptions
-    )
-
-    # Parse options
-    parser = OptionParser(usage=usage)
-    (options, args) = parser.parse_args()
-
-    # Find command
     try:
-        command = args[0]
+        command_name = sys.argv[1]
     except IndexError:
-        parser.print_help()
+        help_index()
         return
 
-    if command in COMMANDS:
-        command = COMMANDS[command]
-        if command.usage is not None:
-            parser.set_usage(command.usage)
-        command.run(parser, options, args)
-    else:
-        parser.error("Unrecognised command: " + command)
+    if command_name == 'help':
+        try:
+            help_command_name = sys.argv[2]
+        except IndexError:
+            help_index()
+            return
+
+        try:
+            command = COMMANDS[help_command_name]
+        except KeyError:
+            unknown_command(help_command_name)
+            return
+
+        command.print_help(help_command_name)
+        return
+
+    try:
+        command = COMMANDS[command_name]
+    except KeyError:
+        unknown_command(command_name)
+        return
+
+    command.execute(sys.argv)
 
 
 if __name__ == "__main__":