Organize subcommands

CHANGES.md

@@ -2,6 +2,19 @@
 
 ## __NEXT__
 
+### Major Changes
+
+* Moved the following modules to subpackages [#1002][]. (@joverlee521)
+  These are technically breaking changes for the API, but they do not change the Augur CLI commands.
+    * `import.py` -> `import_/__init__.py`
+    * `import_beast.py` -> `import_/beast.py`
+    * `measurements.py` -> `measurements/__init__.py` + `measurements/concat.py` + `measurements/export.py`
+* Move the following internal functions/classes [#1002][]. (@joverlee521)
+    * `augur.add_default_command` -> `argparse_.add_default_command`
+    * `utils.HideAsFalseAction` -> `argparse_.HideAsFalseAction`
+* Subcommands must include a `register_parser` function to add their own parser instead of a `register_arguments` function [#1002][]. (@joverlee521)
+
+[#1002]: https://github.com/nextstrain/augur/pull/1002
 
 ## 16.0.3 (6 July 2022)
 

augur/__init__.py

@@ -3,7 +3,6 @@
 """
 
 import argparse
-import re
 import os
 import sys
 import importlib
@@ -13,7 +12,7 @@
 
 from .errors import AugurError
 from .io import print_err
-from .utils import first_line
+from .argparse_ import add_command_subparsers, add_default_command
 
 recursion_limit = os.environ.get("AUGUR_RECURSION_LIMIT")
 if recursion_limit:
@@ -40,7 +39,7 @@
     "export",
     "validate",
     "version",
-    "import",
+    "import_",
     "measurements",
 ]
 
@@ -51,26 +50,11 @@ def make_parser():
         prog        = "augur",
         description = "Augur: A bioinformatics toolkit for phylogenetic analysis.")
 
-    subparsers = parser.add_subparsers()
-
     add_default_command(parser)
     add_version_alias(parser)
 
-    for command in COMMANDS:
-        # Add a subparser for each command.
-        subparser = subparsers.add_parser(
-            command_name(command),
-            help        = first_line(command.__doc__),
-            description = command.__doc__)
-
-        subparser.set_defaults(__command__ = command)
-
-        # Let the command register arguments on its subparser.
-        command.register_arguments(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 = argparse.ArgumentDefaultsHelpFormatter
+    subparsers = parser.add_subparsers()
+    add_command_subparsers(subparsers, COMMANDS)
 
     return parser
 
@@ -99,18 +83,6 @@ def run(argv):
         sys.exit(2)
 
 
-def add_default_command(parser):
-    """
-    Sets the default command to run when none is provided.
-    """
-    class default_command():
-        def run(args):
-            parser.print_help()
-            return 2
-
-    parser.set_defaults(__command__ = default_command)
-
-
 def add_version_alias(parser):
     """
     Add --version as a (hidden) alias for the version command.
@@ -129,17 +101,3 @@ def __call__(self, *args, **kwargs):
         nargs  = 0,
         help   = argparse.SUPPRESS,
         action = run_version_command)
-
-
-def command_name(command):
-    """
-    Returns a short name for a command module.
-    """
-
-    def remove_prefix(prefix, string):
-        return re.sub('^' + re.escape(prefix), '', string)
-
-    package     = command.__package__
-    module_name = command.__name__
-
-    return remove_prefix(package, module_name).lstrip(".").replace("_", "-")

augur/align.py

@@ -17,6 +17,11 @@ class AlignmentError(Exception):
     pass
 
 def register_arguments(parser):
+    """
+    Add arguments to parser.
+    Kept as a separate function than `register_parser` to continue to support
+    unit tests that use this function to create argparser.
+    """
     parser.add_argument('--sequences', '-s', required=True, nargs="+", metavar="FASTA", help="sequences to align")
     parser.add_argument('--output', '-o', default="alignment.fasta", help="output file (default: %(default)s)")
     parser.add_argument('--nthreads', type=nthreads_value, default=1,
@@ -29,6 +34,13 @@ def register_arguments(parser):
     parser.add_argument('--existing-alignment', metavar="FASTA", default=False, help="An existing alignment to which the sequences will be added. The ouput alignment will be the same length as this existing alignment.")
     parser.add_argument('--debug', action="store_true", default=False, help="Produce extra files (e.g. pre- and post-aligner files) which can help with debugging poor alignments.")
 
+
+def register_parser(parent_subparsers):
+    parser = parent_subparsers.add_parser("align", help=__doc__)
+    register_arguments(parser)
+    return parser
+
+
 def prepare(sequences, existing_aln_fname, output, ref_name, ref_seq_fname):
     """Prepare the sequences, existing alignment, and reference sequence for alignment.
 

augur/ancestral.py

@@ -117,7 +117,8 @@ def collect_mutations_and_sequences(tt, infer_tips=False, full_sequences=False,
     return {"nodes": data, "mask": mask}
 
 
-def register_arguments(parser):
+def register_parser(parent_subparsers):
+    parser = parent_subparsers.add_parser("ancestral", help=__doc__)
     parser.add_argument('--tree', '-t', required=True, help="prebuilt Newick")
     parser.add_argument('--alignment', '-a', help="alignment in fasta or VCF format")
     parser.add_argument('--output-node-data', type=str, help='name of JSON file to save mutations and ancestral sequences to')
@@ -133,6 +134,7 @@ def register_arguments(parser):
                                 help='infer nucleotides at ambiguous (N,W,R,..) sites on tip sequences and replace with most likely state.')
     parser.add_argument('--keep-overhangs', action="store_true", default=False,
                                 help='do not infer nucleotides for gaps (-) on either side of the alignment')
+    return parser
 
 def run(args):
     # check alignment type, set flags, read in if VCF

augur/argparse_.py

@@ -0,0 +1,59 @@
+"""
+Custom helpers for the argparse standard library.
+"""
+from argparse import Action, ArgumentDefaultsHelpFormatter
+
+
+def add_default_command(parser):
+    """
+    Sets the default command to run when none is provided.
+    """
+    class default_command():
+        def run(args):
+            parser.print_help()
+            return 2
+
+    parser.set_defaults(__command__ = default_command)
+
+
+def add_command_subparsers(subparsers, commands):
+    """
+    Add subparsers for each command module.
+
+    Parameters
+    ----------
+    subparsers: argparse._SubParsersAction
+        The special subparsers action object created by the parent parser
+        via `parser.add_subparsers()`.
+
+    commands: list[ModuleType]
+        A list of modules that are commands that require their own subparser.
+        Each module is required to have a `register_parser` function to add its own
+        subparser and arguments.
+    """
+    for command in commands:
+        # Allow each command to register its own subparser
+        subparser = command.register_parser(subparsers)
+
+        # Allows us to run commands directly with `args.__command__.run()`
+        subparser.set_defaults(__command__ = command)
+
+        # 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 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)
+
+
+class HideAsFalseAction(Action):
+    """
+    Custom argparse Action that stores False for arguments passed as `--hide*`
+    and stores True for all other argument patterns.
+    """
+    def __call__(self, parser, namespace, values, option_string=None):
+        setattr(namespace, self.dest, option_string[2:6] != 'hide')

augur/clades.py

@@ -248,12 +248,14 @@ def get_reference_sequence_from_root_node(all_muts, root_name):
     return ref
 
 
-def register_arguments(parser):
+def register_parser(parent_subparsers):
+    parser = parent_subparsers.add_parser("clades", help=__doc__)
     parser.add_argument('--tree', help="prebuilt Newick -- no tree will be built if provided")
     parser.add_argument('--mutations', nargs='+', help='JSON(s) containing ancestral and tip nucleotide and/or amino-acid mutations ')
     parser.add_argument('--reference', nargs='+', help='fasta files containing reference and tip nucleotide and/or amino-acid sequences ')
     parser.add_argument('--clades', type=str, help='TSV file containing clade definitions by amino-acid')
     parser.add_argument('--output-node-data', type=str, help='name of JSON file to save clade assignments to')
+    return parser
 
 
 def run(args):

augur/distance.py

@@ -159,7 +159,7 @@
 
 from .frequency_estimators import timestamp_to_float
 from .reconstruct_sequences import load_alignments
-from .utils import annotate_parents_for_tree, read_node_data, write_json
+from .utils import annotate_parents_for_tree, first_line, read_node_data, write_json
 
 
 def read_distance_map(map_file):
@@ -626,7 +626,8 @@ def get_distances_to_all_pairs(tree, sequences_by_node_and_gene, distance_map, e
     return distances_by_node
 
 
-def register_arguments(parser):
+def register_parser(parent_subparsers):
+    parser = parent_subparsers.add_parser("distance", help=first_line(__doc__))
     parser.add_argument("--tree", help="Newick tree", required=True)
     parser.add_argument("--alignment", nargs="+", help="sequence(s) to be used, supplied as FASTA files", required=True)
     parser.add_argument('--gene-names', nargs="+", type=str, help="names of the sequences in the alignment, same order assumed", required=True)
@@ -637,6 +638,7 @@ def register_arguments(parser):
     parser.add_argument("--earliest-date", help="earliest date at which samples are considered to be from previous seasons (e.g., 2019-01-01). This date is only used in pairwise comparisons. If omitted, all samples prior to the latest date will be considered.")
     parser.add_argument("--latest-date", help="latest date at which samples are considered to be from previous seasons (e.g., 2019-01-01); samples from any date after this are considered part of the current season")
     parser.add_argument("--output", help="JSON file with calculated distances stored by node name and attribute name", required=True)
+    return parser
 
 
 def run(args):

augur/export.py

@@ -1,22 +1,21 @@
 """
 Export JSON files suitable for visualization with auspice.
 """
-from .export_v1 import run_v1, register_arguments_v1
-from .export_v2 import run_v2, register_arguments_v2
+from .argparse_ import add_command_subparsers
+from . import export_v1, export_v2
 
+SUBCOMMANDS = [
+    export_v1,
+    export_v2,
+]
 
-def register_arguments(parser):
+
+def register_parser(parent_subparsers):
+    parser = parent_subparsers.add_parser("export", help=__doc__)
+    # Add subparsers for subcommands
     metavar_msg ="Augur export now needs you to define the JSON version " + \
                  "you want, e.g. `augur export v2`."
     subparsers = parser.add_subparsers(title="JSON SCHEMA",
                                        metavar=metavar_msg)
-    subparsers.required = True
-    register_arguments_v2(subparsers)
-    register_arguments_v1(subparsers)
-
-
-def run(args):
-    if "v1" in args:
-        return run_v1(args)
-    else:
-        return run_v2(args)
+    add_command_subparsers(subparsers, SUBCOMMANDS)
+    return parser

augur/export_v1.py

@@ -1,6 +1,5 @@
 """
-Augur Version 1-specific helpers for exporting JSON files suitable for
-visualization with auspice.
+Export version 1 JSON schema (separate meta and tree JSONs) for visualization with Auspice
 """
 
 import os, sys
@@ -261,7 +260,7 @@ def add_tsv_metadata_to_nodes(nodes, meta_tsv, meta_json, extra_fields=['authors
         fields = [x for x in meta_json["color_options"].keys() if x != "gt"] + extra_fields
     else:
         fields = list(extra_fields)
-        
+
     if "geo" in meta_json:
         fields += meta_json["geo"]
 
@@ -331,15 +330,15 @@ def add_option_args(parser):
     return options
 
 
-def register_arguments_v1(subparsers):
-    # V1 sub-command
-    v1 = subparsers.add_parser('v1', help="Export version 1 JSON schema (separate meta and tree JSONs)")
-    v1_core = add_core_args(v1)
-    v1_options = add_option_args(v1)
-    v1.add_argument("--v1", help=SUPPRESS, default=True)
+def register_parser(parent_subparsers):
+    parser = parent_subparsers.add_parser("v1", help=__doc__)
+    add_core_args(parser)
+    add_option_args(parser)
+    parser.add_argument("--v1", help=SUPPRESS, default=True)
+    return parser
 
 
-def run_v1(args):
+def run(args):
     T = Phylo.read(args.tree, 'newick')
     node_data = read_node_data(args.node_data) # args.node_data is an array of multiple files (or a single file)
     nodes = node_data["nodes"] # this is the per-node metadata produced by various augur modules

augur/export_v2.py

@@ -1,5 +1,5 @@
 """
-Export JSON files suitable for visualization with auspice.
+Export version 2 JSON schema for visualization with Auspice
 """
 from pathlib import Path
 import os, sys
@@ -819,17 +819,17 @@ def node_data_prop_is_normal_trait(name):
     return True
 
 
-def register_arguments_v2(subparsers):
-    v2 = subparsers.add_parser("v2", help="Export version 2 JSON schema")
+def register_parser(parent_subparsers):
+    parser = parent_subparsers.add_parser("v2", help=__doc__)
 
-    required = v2.add_argument_group(
+    required = parser.add_argument_group(
         title="REQUIRED"
     )
     required.add_argument('--tree','-t', metavar="newick", required=True, help="Phylogenetic tree, usually output from `augur refine`")
     required.add_argument('--node-data', metavar="JSON", required=True, nargs='+', help="JSON files containing metadata for nodes in the tree")
     required.add_argument('--output', metavar="JSON", required=True, help="Ouput file (typically for visualisation in auspice)")
 
-    config = v2.add_argument_group(
+    config = parser.add_argument_group(
         title="DISPLAY CONFIGURATION",
         description="These control the display settings for auspice. \
             You can supply a config JSON (which has all available options) or command line arguments (which are more limited but great to get started). \
@@ -844,21 +844,21 @@ def register_arguments_v2(subparsers):
     config.add_argument('--color-by-metadata', metavar="trait", nargs='+', help="Metadata columns to include as coloring options")
     config.add_argument('--panels', metavar="panels", nargs='+', choices=['tree', 'map', 'entropy', 'frequencies', 'measurements'], help="Restrict panel display in auspice. Options are %(choices)s. Ignore this option to display all available panels.")
 
-    optional_inputs = v2.add_argument_group(
+    optional_inputs = parser.add_argument_group(
         title="OPTIONAL INPUT FILES"
     )
     optional_inputs.add_argument('--metadata', metavar="FILE", help="Additional metadata for strains in the tree, as CSV or TSV")
     optional_inputs.add_argument('--colors', metavar="FILE", help="Custom color definitions, one per line in the format `TRAIT_TYPE\\tTRAIT_VALUE\\tHEX_CODE`")
     optional_inputs.add_argument('--lat-longs', metavar="TSV", help="Latitudes and longitudes for geography traits (overrides built in mappings)")
 
-    optional_settings = v2.add_argument_group(
+    optional_settings = parser.add_argument_group(
         title="OPTIONAL SETTINGS"
     )
     optional_settings.add_argument('--minify-json', action="store_true", help="export JSONs without indentation or line returns")
     optional_settings.add_argument('--include-root-sequence', action="store_true", help="Export an additional JSON containing the root sequence (reference sequence for vcf) used to identify mutations. The filename will follow the pattern of <OUTPUT>_root-sequence.json for a main auspice JSON of <OUTPUT>.json")
     optional_settings.add_argument('--skip-validation', action="store_true", help="skip validation of input/output files. Use at your own risk!")
 
-    return v2
+    return parser
 
 
 def set_display_defaults(data_json, config):
@@ -982,7 +982,7 @@ def get_config(args):
         del config["vaccine_choices"]
     return config
 
-def run_v2(args):
+def run(args):
     configure_warnings()
     data_json = {"version": "v2", "meta": {"updated": time.strftime('%Y-%m-%d')}}