minor refactoring

metadex/__main__.py

@@ -3,6 +3,7 @@ import logging
 import os
 import stat
 import sys
+from contextlib import contextmanager
 from pathlib import Path
 
 from . import config, metadex, utils
@@ -10,6 +11,18 @@ from . import config, metadex, utils
 log = logging.getLogger(__name__)
 
 
+_commands = {}
+
+
+def command(name):
+    def wrapper(f):
+        assert name not in _commands
+        _commands[name] = f
+        return f
+
+    return wrapper
+
+
 def getargs():
     parser = argparse.ArgumentParser()
     parser.set_defaults(mode=None)
@@ -41,114 +54,116 @@ def getargs():
 
     subparsers = parser.add_subparsers(title="commands")
 
+    @contextmanager
+    def command_parser(name, **kwargs):
+        assert name in _commands
+        subparser = subparsers.add_parser(name, **kwargs)
+        subparser.set_defaults(mode=name)
+        yield subparser
+
     # Command: scan
 
-    parser_scan = subparsers.add_parser("scan", help="scan a local file system")
-    parser_scan.set_defaults(mode="scan")
+    with command_parser("scan", help="scan a local file system") as subparser:
 
-    parser_scan.add_argument(
-        "basedir",
-        type=Path,
-        nargs="+",
-        help="index all files from this dir",
-    )
-    parser_scan.add_argument(
-        "--no-remove-missing",
-        dest="remove_missing",
-        action="store_false",
-        help="do not remove files from the database which cannot be found in the file system",
-    )
-    parser_scan.add_argument(
-        "--map-mount",
-        nargs="*",
-        default=[],
-        type=str,
-        help="map a source host:path to any other destination during scanning for files\nExample: src=/mnt/foo,dest=foo:",
-    )
+        subparser.add_argument(
+            "basedir",
+            type=Path,
+            nargs="+",
+            help="index all files from this dir",
+        )
+        subparser.add_argument(
+            "--no-remove-missing",
+            dest="remove_missing",
+            action="store_false",
+            help="do not remove files from the database which cannot be found in the file system",
+        )
+        subparser.add_argument(
+            "--map-mount",
+            nargs="*",
+            default=[],
+            type=str,
+            help="map a source host:path to any other destination during scanning for files\nExample: src=/mnt/foo,dest=foo:",
+        )
 
     # Command: ingest-ls
 
-    parser_ingest_ls = subparsers.add_parser(
+    with command_parser(
         "ingest-ls",
         help="ingest extra data",
         description="When ingesting data from an external source, the hostname will not be set automatically.",
-    )
-    parser_ingest_ls.set_defaults(mode="ingest-ls")
+    ) as subparser:
 
-    parser_ingest_ls.add_argument(
-        "infile",
-        nargs="?",
-        type=argparse.FileType(),
-        default=sys.stdin,
-        help="output from `ls -lR`",
-    )
-    parser_ingest_ls.add_argument(
-        "--remove-missing",
-        action="store_true",
-        help="Remove files not listed in the infile.",
-    )
-    parser_ingest_ls.add_argument(
-        "--ref-year",
-        type=int,
-        help="The year when 'ls -l' was run, to resolve relative dates.",
-    )
+        subparser.add_argument(
+            "infile",
+            nargs="?",
+            type=argparse.FileType(),
+            default=sys.stdin,
+            help="output from `ls -lR`",
+        )
+        subparser.add_argument(
+            "--remove-missing",
+            action="store_true",
+            help="Remove files not listed in the infile.",
+        )
+        subparser.add_argument(
+            "--ref-year",
+            type=int,
+            help="The year when 'ls -l' was run, to resolve relative dates.",
+        )
 
     # Command: ingest-db
 
-    parser_ingest_db = subparsers.add_parser("ingest-db")
-    parser_ingest_db.set_defaults(mode="ingest-db")
+    with command_parser("ingest-db") as subparser:
 
-    parser_ingest_db.add_argument(
-        "infile",
-        type=Path,
-        help="a Metadex SQLite DB",
-    )
-    parser_ingest_db.add_argument(
-        "--map-mount",
-        nargs="*",
-        default=[],
-        type=str,
-        help="map a source host:path to any other destination while importing",
-    )
+        subparser.add_argument(
+            "infile",
+            type=Path,
+            help="a Metadex SQLite DB",
+        )
+        subparser.add_argument(
+            "--map-mount",
+            nargs="*",
+            default=[],
+            type=str,
+            help="map a source host:path to any other destination while importing",
+        )
 
     # Command: rm
 
-    parser_rm = subparsers.add_parser("rm")
-    parser_rm.set_defaults(mode="rm")
-    parser_rm.add_argument(
-        "files",
-        type=str,
-        nargs="+",
-        help="files to remove",
-    )
-    parser_rm.add_argument(
-        "-r",
-        action="store_true",
-        dest="include_subfiles",
-        help="include sub-files",
-    )
+    with command_parser("rm") as subparser:
+        subparser.add_argument(
+            "files",
+            type=str,
+            nargs="+",
+            help="files to remove",
+        )
+        subparser.add_argument(
+            "-r",
+            action="store_true",
+            dest="include_subfiles",
+            help="include sub-files",
+        )
 
     # Command: ls
 
-    parser_ls = subparsers.add_parser("ls")
-    parser_ls.set_defaults(mode="ls")
-    parser_ls.add_argument(
-        "file",
-        type=str,
-        nargs="*",
-        help="look up a file",
-    )
-    parser_ls.add_argument(
-        "--type",
-        "-t",
-        choices="dfl",
-        help="Filter searches to (d)irectories, plain (f)iles, or sym(l)inks.",
-    )
-    parser_ls.add_argument("--format", type=str, default="{date}\t{size}\t{path}")
-    parser_ls.add_argument(
-        "--match", choices=("regex", "glob", "fuzzy"), default="glob"
-    )
-    parser_ls.add_argument("--stop-on-error", action="store_true")
+    with command_parser("ls") as subparser:
+        subparser.add_argument(
+            "file",
+            type=str,
+            nargs="*",
+            help="look up a file",
+        )
+        subparser.add_argument(
+            "--type",
+            "-t",
+            choices="dfl",
+            help="Filter searches to (d)irectories, plain (f)iles, or sym(l)inks.",
+        )
+        subparser.add_argument("--format", type=str, default="{date}\t{size}\t{path}")
+        subparser.add_argument(
+            "--match", choices=("regex", "glob", "fuzzy"), default="glob"
+        )
+        subparser.add_argument("--stop-on-error", action="store_true")
 
     # Parse args.
 
@@ -167,6 +182,7 @@ def getargs():
     return args
 
 
+@command("ingest-ls")
 def cmd_ingest_ls(args):
     metadex.init(args.db)
 
@@ -181,6 +197,7 @@ def cmd_ingest_ls(args):
     metadex.close()
 
 
+@command("ingest-db")
 def cmd_ingest_db(args):
     metadex.init(args.db)
 
@@ -195,6 +212,7 @@ def cmd_ingest_db(args):
     metadex.close()
 
 
+@command("scan")
 def cmd_scan(args):
     metadex.init(args.db)
 
@@ -213,6 +231,7 @@ def cmd_scan(args):
     metadex.close()
 
 
+@command("rm")
 def cmd_rm(args):
     metadex.init(args.db)
 
@@ -221,6 +240,7 @@ def cmd_rm(args):
 
     metadex.close()
 
+@command("ls")
 def cmd_ls(args) -> int:
     return_code = 0
 
@@ -316,19 +336,12 @@ def main():
     if config.dryrun:
         log.info(f"--- DRY RUN ---")
 
-    if args.mode == "scan":
-        return cmd_scan(args)
-    elif args.mode == "ingest-ls":
-        return cmd_ingest_ls(args)
-    elif args.mode == "ingest-db":
-        return cmd_ingest_db(args)
-    elif args.mode == "rm":
-        return cmd_rm(args)
-    elif args.mode == "ls":
+    if args.mode == "ls":
         # Since this is a read-only operation we can change some config params.
         config.db_allow_slow = False
         config.dryrun = True
-        return cmd_ls(args)
+    cmd = _commands[args.mode]
+    return cmd(args)
 
 
 if __name__ == "__main__":