🚧 Register command parsers lazily

This isn't in a working state yet, but augur -h does run without
importing every subcommand module which is a speed improvement.

Lots of rough edges:

1. Command names are taken literally, so `augur import` needs to be
   `augur import_`. This can be fixed with a predefined mapping.
2. Top-level commands (e.g. `augur -h`) run, but help text is missing.
3. First-level subcommands (e.g. `augur filter -h`) run, but all options
   are missing.
4. Second-level subcommands (e.g. `augur import beast`) do not work yet.

Inspired by https://stackoverflow.com/a/51903799

augur/__init__.py

@@ -5,7 +5,6 @@
 import argparse
 import os
 import sys
-import importlib
 import traceback
 from textwrap import dedent
 from types import SimpleNamespace
@@ -48,18 +47,28 @@
     "write_file",
 ]
 
-COMMANDS = [importlib.import_module('augur.' + c) for c in command_strings]
+class LazyArgumentParser(argparse.ArgumentParser):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.lazy_init = None
+
+    def parse_known_args(self, args=None, namespace=None):
+        if self.lazy_init is not None:
+            self.lazy_init()
+            self.lazy_init = None
+        return super().parse_known_args(args, namespace)
+
 
 def make_parser():
-    parser = argparse.ArgumentParser(
+    parser = LazyArgumentParser(
         prog        = "augur",
         description = "Augur: A bioinformatics toolkit for phylogenetic analysis.")
 
     add_default_command(parser)
     add_version_alias(parser)
 
     subparsers = parser.add_subparsers()
-    add_command_subparsers(subparsers, COMMANDS)
+    add_command_subparsers(subparsers, command_strings)
 
     return parser
 

augur/argparse_.py

@@ -2,8 +2,10 @@
 Custom helpers for the argparse standard library.
 """
 from argparse import Action, ArgumentDefaultsHelpFormatter, ArgumentParser, _ArgumentGroup
+import importlib
 from typing import Union
 from .types import ValidationMode
+from augur.utils import first_line
 
 
 # Include this in an argument help string to suppress the automatic appending
@@ -32,7 +34,7 @@ def run(args):
     parser.set_defaults(__command__ = default_command)
 
 
-def add_command_subparsers(subparsers, commands, command_attribute='__command__'):
+def add_command_subparsers(subparsers, command_strings, command_attribute='__command__'):
     """
     Add subparsers for each command module.
 
@@ -51,24 +53,34 @@ def add_command_subparsers(subparsers, commands, command_attribute='__command__'
         Optional attribute name for the commands. The default is `__command__`,
         which allows top level augur to run commands directly via `args.__command__.run()`.
     """
-    for command in commands:
-        # Allow each command to register its own subparser
-        subparser = command.register_parser(subparsers)
+    for command_string in command_strings:
+        # FIXME: add help=first_line(command.__doc__). not trivial because command is not yet imported
+        subparser = subparsers.add_parser(command_string)
 
-        # Add default attribute for command module
-        if command_attribute:
-            subparser.set_defaults(**{command_attribute: command})
+        def make_lazy_init(command_string):
+            def lazy_init():
+                command = importlib.import_module('augur.' + command_string)
+                # Allow each command to register its own subparser
+                command.register_parser(subparser)
 
-        # Use the same formatting class for every command for consistency.
-        # Set here to avoid repeating it in every command's register_parser().
-        subparser.formatter_class = ArgumentDefaultsHelpFormatter
+                # Add default attribute for command module
+                if command_attribute:
+                    subparser.set_defaults(**{command_attribute: command})
 
-        if not subparser.description and command.__doc__:
-            subparser.description = command.__doc__
+                # Use the same formatting class for every command for consistency.
+                # Set here to avoid repeating it in every command's register_parser().
+                subparser.formatter_class = ArgumentDefaultsHelpFormatter
 
-        # If a command doesn't have its own run() function, then print its help when called.
-        if not getattr(command, "run", None):
-            add_default_command(subparser)
+                if not subparser.description and command.__doc__:
+                    subparser.description = command.__doc__
+
+                # If a command doesn't have its own run() function, then print its help when called.
+                if not getattr(command, "run", None):
+                    add_default_command(subparser)
+
+            return lazy_init
+
+        subparser.lazy_init = make_lazy_init(command_string)
 
 
 class HideAsFalseAction(Action):

augur/filter/__init__.py

@@ -117,8 +117,7 @@ def register_arguments(parser):
     parser.set_defaults(probabilistic_sampling=True)
 
 
-def register_parser(parent_subparsers):
-    parser = parent_subparsers.add_parser("filter", help=__doc__)
+def register_parser(parser):
     register_arguments(parser)
     return parser
 

augur/import_/__init__.py

@@ -2,15 +2,12 @@
 Import analyses into augur pipeline from other systems
 """
 from augur.argparse_ import add_command_subparsers
-from augur.utils import first_line
-from . import beast
 
 SUBCOMMANDS = [
-    beast,
+    "import_.beast",
 ]
 
-def register_parser(parent_subparsers):
-    parser = parent_subparsers.add_parser("import", help=first_line(__doc__))
+def register_parser(parser):
     metavar_msg = "Import analyses into augur pipeline from other systems"
     subparsers = parser.add_subparsers(title="TYPE",
                                        metavar=metavar_msg)