Merge pull request #16 from hajs/patch-1

Author: casperdcl

shtab/main.py

@@ -24,6 +24,9 @@ def get_main_parser():
         "--prefix", help="prepended to generated functions to avoid clashes"
     )
     parser.add_argument("--preamble", help="prepended to generated script")
+    parser.add_argument(
+        "--prog", help="custom program name (overrides `parser.prog`)"
+    )
     parser.add_argument(
         "-u",
         "--error-unimportable",
@@ -50,6 +53,8 @@ def main(argv=None):
     other_parser = getattr(module, other_parser)
     if callable(other_parser):
         other_parser = other_parser()
+    if args.prog:
+        other_parser.prog = args.prog
     print(
         complete(
             other_parser,

tests/test_shtab.py

@@ -14,6 +14,7 @@
 from shtab.main import get_main_parser, main
 
 SUPPORTED_SHELLS = "bash", "zsh"
+fix_shell = pytest.mark.parametrize("shell", SUPPORTED_SHELLS)
 
 
 class Bash(object):
@@ -70,15 +71,42 @@ def test_choices():
     assert "" not in shtab.Required.FILE
 
 
-@pytest.mark.parametrize("shell", SUPPORTED_SHELLS)
+@fix_shell
 def test_main(shell, caplog):
     with caplog.at_level(logging.INFO):
         main(["-s", shell, "shtab.main.get_main_parser"])
 
     assert not caplog.record_tuples
 
 
-@pytest.mark.parametrize("shell", SUPPORTED_SHELLS)
+@fix_shell
+def test_prog_override(shell, caplog, capsys):
+    with caplog.at_level(logging.INFO):
+        main(["-s", shell, "--prog", "foo", "shtab.main.get_main_parser"])
+
+    captured = capsys.readouterr()
+    assert not captured.err
+    if shell == "bash":
+        assert "complete -o nospace -F _shtab_shtab foo" in captured.out
+
+    assert not caplog.record_tuples
+
+
+@fix_shell
+def test_prefix_override(shell, caplog, capsys):
+    with caplog.at_level(logging.INFO):
+        main(["-s", shell, "--prefix", "foo", "shtab.main.get_main_parser"])
+
+    captured = capsys.readouterr()
+    assert not captured.err
+    if shell == "bash":
+        shell = Bash(captured.out)
+        shell.compgen('-W "$_shtab_foo_options_"', "--h", "--help")
+
+    assert not caplog.record_tuples
+
+
+@fix_shell
 def test_complete(shell, caplog):
     parser = get_main_parser()
     with caplog.at_level(logging.INFO):
@@ -92,7 +120,7 @@ def test_complete(shell, caplog):
     assert not caplog.record_tuples
 
 
-@pytest.mark.parametrize("shell", SUPPORTED_SHELLS)
+@fix_shell
 def test_positional_choices(shell, caplog):
     parser = ArgumentParser(prog="test")
     parser.add_argument("posA", choices=["one", "two"])