Restructure how subcommands work

flit/__init__.py

@@ -1,92 +1,46 @@
 """A simple packaging tool for simple packages."""
 import argparse
 import logging
-import pathlib
+from pathlib import Path
 import sys
 
 from . import common
 from .log import enable_colourful_output
+from .subcmds import Subcommand, SubcommandArgumentParser
 
 __version__ = '0.13'
 
 log = logging.getLogger(__name__)
 
-def add_shared_install_options(parser):
-    parser.add_argument('--user', action='store_true', default=None,
-        help="Do a user-local install (default if site.ENABLE_USER_SITE is True)"
-    )
-    parser.add_argument('--env', action='store_false', dest='user',
-        help="Install into sys.prefix (default if site.ENABLE_USER_SITE is False, i.e. in virtualenvs)"
-    )
-    parser.add_argument('--python', default=sys.executable,
-        help="Target Python executable, if different from the one running flit"
+def add_ini_file_option(parser):
+    default = pyproject = Path('pyproject.toml')
+    flit_ini = Path('flit.ini')
+    if flit_ini.is_file() and not pyproject.is_file():
+        default = flit_ini
+    parser.add_argument('-f', '--ini-file', type=Path, default=default,
+        help=""
     )
 
+
+subcmds = [
+    Subcommand('build', func='flit.build:main', help="Build wheel and sdist"),
+    Subcommand('publish', func='flit.upload:main', help="Upload wheel and sdist"),
+    Subcommand('install', func='flit.install:main', help="Install the package"),
+    Subcommand('installfrom', func='flit.installfrom:main',
+               help="Download and install a package using flit from source"),
+    Subcommand('init', func='flit.init:main',
+               help="Prepare pyproject.toml for a new package")
+]
+
 def main(argv=None):
-    ap = argparse.ArgumentParser()
-    ap.add_argument('-f', '--ini-file', type=pathlib.Path, default='pyproject.toml')
+    ap = SubcommandArgumentParser()
     ap.add_argument('--version', action='version', version='Flit '+__version__)
-    ap.add_argument('--repository',
-        help="Name of the repository to upload to (must be in ~/.pypirc)"
-    )
     ap.add_argument('--debug', action='store_true', help=argparse.SUPPRESS)
     ap.add_argument('--logo', action='store_true', help=argparse.SUPPRESS)
-    subparsers = ap.add_subparsers(title='subcommands', dest='subcmd')
-
-    parser_build = subparsers.add_parser('build',
-        help="Build wheel and sdist",
-    )
-
-    parser_build.add_argument('--format', action='append',
-        help="Select a format to build. Options: 'wheel', 'sdist'"
-    )
-
-    parser_publish = subparsers.add_parser('publish',
-        help="Upload wheel and sdist",
-    )
-
-    parser_publish.add_argument('--format', action='append',
-        help="Select a format to publish. Options: 'wheel', 'sdist'"
-    )
-
-    parser_install = subparsers.add_parser('install',
-        help="Install the package",
-    )
-    parser_install.add_argument('-s', '--symlink', action='store_true',
-        help="Symlink the module/package into site packages instead of copying it"
-    )
-    parser_install.add_argument('--pth-file', action='store_true',
-        help="Add .pth file for the module/package to site packages instead of copying it"
-    )
-    add_shared_install_options(parser_install)
-    parser_install.add_argument('--deps', choices=['all', 'production', 'develop', 'none'], default='all',
-        help="Which set of dependencies to install")
-
-    parser_installfrom = subparsers.add_parser('installfrom',
-       help="Download and install a package using flit from source"
-    )
-    parser_installfrom.add_argument('location',
-        help="A URL to download, or a shorthand like github:takluyver/flit"
-    )
-    add_shared_install_options(parser_installfrom)
-
-    parser_init = subparsers.add_parser('init',
-        help="Prepare flit.ini for a new package"
-    )
+    ap.add_subcommands(subcmds)
 
     args = ap.parse_args(argv)
 
-    cf = args.ini_file
-    if args.subcmd != 'init' and cf == pathlib.Path('pyproject.toml')\
-            and not cf.is_file():
-        # Fallback to flit.ini if it's present
-        cf_ini = pathlib.Path('flit.ini')
-        if cf_ini.is_file():
-            args.ini_file = cf_ini
-        else:
-            sys.exit('Neither pyproject.toml nor flit.ini found, '
-                     'and no other config file path specified')
-
     enable_colourful_output(logging.DEBUG if args.debug else logging.INFO)
 
     log.debug("Parsed arguments %r", args)
@@ -96,26 +50,4 @@ def main(argv=None):
         print(clogo.format(version=__version__))
         sys.exit(0)
 
-    if args.subcmd == 'build':
-        from .build import main
-        main(args.ini_file, formats=set(args.format or []))
-    elif args.subcmd == 'publish':
-        from .upload import main
-        main(args.ini_file, args.repository, formats=set(args.format or []))
-
-    elif args.subcmd == 'install':
-        from .install import Installer
-        try:
-            Installer(args.ini_file, user=args.user, python=args.python,
-                      symlink=args.symlink, deps=args.deps, pth=args.pth_file).install()
-        except (common.NoDocstringError, common.NoVersionError) as e:
-            sys.exit(e.args[0])
-    elif args.subcmd == 'installfrom':
-        from .installfrom import installfrom
-        sys.exit(installfrom(args.location, user=args.user, python=args.python))
-    elif args.subcmd == 'init':
-        from .init import TerminalIniter
-        TerminalIniter().initialise()
-    else:
-        ap.print_help()
-        sys.exit(1)
+    ap.dispatch_subcommand(args)

flit/build.py

@@ -1,5 +1,6 @@
 """flit build - build both wheel and sdist"""
 
+import argparse
 from contextlib import contextmanager
 import logging
 import os
@@ -26,7 +27,7 @@ def unpacked_tarball(path):
         assert len(files) == 1, files
         yield os.path.join(tmpdir, files[0])
 
-def main(ini_file: Path, formats=None):
+def build(ini_file: Path, formats=None):
     """Build wheel and sdist"""
     if not formats:
         formats = ALL_FORMATS
@@ -54,3 +55,20 @@ def main(ini_file: Path, formats=None):
         sys.exit('Config error: {}'.format(e))
 
     return SimpleNamespace(wheel=wheel_info, sdist=sdist_info)
+
+
+def add_format_option(parser):
+    parser.add_argument('--format', action='append',
+        help="Select a format to build. Options: 'wheel', 'sdist'"
+    )
+
+
+
+def main(argv):
+    ap = argparse.ArgumentParser(prog='flit build')
+    from . import add_ini_file_option
+    add_ini_file_option(ap)
+    add_format_option(ap)
+    args = ap.parse_args(argv)
+
+    build(args.ini_file, formats=set(args.format or []))

flit/init.py

@@ -198,5 +198,8 @@ def initialise(self):
 {metadata}
 """
 
+def main(argv):
+    TerminalIniter().initialise()
+
 if __name__ == '__main__':
     TerminalIniter().initialise()

flit/install.py

@@ -1,5 +1,6 @@
 """Install packages locally for development
 """
+import argparse
 import logging
 import os
 import csv
@@ -348,3 +349,38 @@ def install(self):
             self.install_directly()
         else:
             self.install_with_pip()
+
+
+def add_shared_install_options(parser):
+    parser.add_argument('--user', action='store_true', default=None,
+        help="Do a user-local install (default if site.ENABLE_USER_SITE is True)"
+    )
+    parser.add_argument('--env', action='store_false', dest='user',
+        help="Install into sys.prefix (default if site.ENABLE_USER_SITE is False, i.e. in virtualenvs)"
+    )
+    parser.add_argument('--python', default=sys.executable,
+        help="Target Python executable, if different from the one running flit"
+    )
+
+def main(argv):
+    ap = argparse.ArgumentParser()
+    from . import add_ini_file_option
+    add_ini_file_option(ap)
+    add_shared_install_options(ap)
+    ap.add_argument('-s', '--symlink', action='store_true',
+        help="Symlink the module/package into site packages instead of copying it"
+    )
+    ap.add_argument('--pth-file', action='store_true',
+        help="Add .pth file for the module/package to site packages instead of copying it"
+    )
+    ap.add_argument('--deps', choices=['all', 'production', 'develop', 'none'],
+        default='all', help="Which set of dependencies to install"
+    )
+    args = ap.parse_args(argv)
+
+    try:
+        Installer(args.ini_file, user=args.user, python=args.python,
+                  symlink=args.symlink, deps=args.deps,
+                  pth=args.pth_file).install()
+    except (common.NoDocstringError, common.NoVersionError) as e:
+        sys.exit(e.args[0])

flit/installfrom.py

@@ -1,3 +1,4 @@
+import argparse
 import os.path
 import pathlib
 import re
@@ -124,3 +125,14 @@ def installfrom(address, user=None, python=sys.executable):
     except BadInput as e:
         print(e, file=sys.stderr)
         return 2
+
+def main(argv):
+    ap = argparse.ArgumentParser()
+    from .install import add_shared_install_options
+    add_shared_install_options(ap)
+    ap.add_argument('location',
+        help="A URL to download, or a shorthand like github:takluyver/flit"
+    )
+    args = ap.parse_args(argv)
+
+    return installfrom(args.location, user=args.user, python=args.python)

flit/subcmds.py

@@ -0,0 +1,78 @@
+import argparse
+import importlib
+import sys
+
+def load_obj(objref):
+    """Load an object from an entry point style object reference
+
+    'mod.submod:obj.attr'
+    """
+    modname, qualname_separator, qualname = objref.partition(':')
+    obj = importlib.import_module(modname)
+    if qualname_separator:
+        for attr in qualname.split('.'):
+            obj = getattr(obj, attr)
+    return obj
+
+class Subcommand:
+    def __init__(self, name, *, func, help):
+        self.name = name
+        self.func = func
+        self.help = help
+
+    def run(self, extra_argv):
+        if isinstance(self.func, str):
+            self.func = load_obj(self.func)
+        return self.func(extra_argv)
+
+class SubcommandArgumentParser(argparse.ArgumentParser):
+    def add_subcommands(self, subcommands):
+        self.add_argument('subcommand_and_args', nargs=argparse.REMAINDER)
+        self.__subcommands = subcommands
+
+
+    def format_help(self):
+        special_options = self._optionals._group_actions
+        special_options_names = []
+        for so in special_options:
+            if so.help is not argparse.SUPPRESS:
+                special_options_names.extend(so.option_strings)
+
+        # usage
+        lines = [
+            "usage: {prog} [{options}]".format(prog=self.prog,
+               options=" | ".join(special_options_names)),
+            "   or: {prog} <command>".format(prog=self.prog),
+            ""
+        ]
+
+        # special optional args
+        formatter = self._get_formatter()
+        formatter.start_section("Special options")
+        formatter.add_arguments(self._optionals._group_actions)
+        formatter.end_section()
+        lines.append(formatter.format_help())
+
+        # subcommands
+        lines += ["Commands:"]
+        for sc in self.__subcommands:
+            if sc.help is not argparse.SUPPRESS:
+                s = "  {:<12} {}".format(sc.name, sc.help)
+                lines.append(s)
+
+        return "\n".join(lines) + "\n"
+
+    def dispatch_subcommand(self, parsed_args):
+        if not parsed_args.subcommand_and_args:
+            self.print_help()
+            sys.exit(2)
+
+        subcmd, *extra_argv = parsed_args.subcommand_and_args
+        for sc in self.__subcommands:
+            if sc.name == subcmd:
+                sys.exit(sc.run(extra_argv))
+
+        print("Unknown command {!r}".format(subcmd))
+        print("  Available commands are: ",
+              ", ".join(sc.name for sc in self.__subcommands))
+        sys.exit(2)

flit/upload.py

@@ -3,6 +3,7 @@
 This is cribbed heavily from distutils.command.(upgrade|register), which as part
 of Python is under the PSF license.
 """
+import argparse
 import configparser
 import getpass
 import hashlib
@@ -269,7 +270,7 @@ def do_upload(file:Path, metadata:Metadata, repo_name=None):
         log.info("Package is at %s/%s", repo['url'], metadata.name)
 
 
-def main(ini_path, repo_name, formats=None):
+def publish(ini_path, repo_name, formats=None):
     """Build and upload wheel and sdist."""
     from . import build
     built = build.main(ini_path, formats=formats)
@@ -278,3 +279,16 @@ def main(ini_path, repo_name, formats=None):
         do_upload(built.wheel.file, built.wheel.builder.metadata, repo_name)
     if built.sdist is not None:
         do_upload(built.sdist.file, built.sdist.builder.metadata, repo_name)
+
+def main(argv):
+    ap = argparse.ArgumentParser()
+    from . import add_ini_file_option
+    add_ini_file_option(ap)
+    from .build import add_format_option
+    add_format_option(ap)
+    ap.add_argument('--repository',
+        help="Name of the repository to upload to (must be in ~/.pypirc)"
+    )
+    args = ap.parse_args(argv)
+
+    publish(args.ini_file, args.repository, formats=set(args.format or []))

tests/test_build.py

@@ -18,7 +18,7 @@
     print('EG_README.rst')
 """.format(python=sys.executable)
 
-def test_build_main():
+def test_build():
     with TemporaryDirectory() as td:
         pyproject = Path(td, 'pyproject.toml')
         shutil.copy(str(samples_dir / 'module1-pkg.toml'), str(pyproject))
@@ -27,7 +27,7 @@ def test_build_main():
         Path(td, '.git').mkdir()   # Fake a git repo
 
         with MockCommand('git', LIST_FILES):
-            res = build.main(pyproject)
+            res = build.build(pyproject)
         assert res.wheel.file.suffix == '.whl'
         assert res.sdist.file.name.endswith('.tar.gz')
 

tests/test_command.py

@@ -10,4 +10,4 @@ def test_flit_usage():
     p = Popen([sys.executable, '-m', 'flit'], stdout=PIPE, stderr=STDOUT)
     out, _ = p.communicate()
     assert 'Build wheel' in out.decode('utf-8', 'replace')
-    assert p.poll() == 1
+    assert p.poll() == 2