misc tidy

Author: casperdcl

shtab/main.py

@@ -1,34 +1,23 @@
 import argparse
-import contextlib
 import logging
 import os
 import sys
+from contextlib import contextmanager
 from importlib import import_module
 from pathlib import Path
-from typing import Generator
-from typing import Optional as Opt
-from typing import TextIO
 
 from . import SUPPORTED_SHELLS, __version__, add_argument_to, complete
 
 log = logging.getLogger(__name__)
 
 
-@contextlib.contextmanager
-def extract_stdout(output: Opt[Path]) -> Generator[TextIO, None, None]:
-    if output is None:
-        yield sys.stdout
-    else:
-        with output.open("w") as stdout:
-            yield stdout
-
-
 def get_main_parser():
     parser = argparse.ArgumentParser(prog="shtab")
     parser.add_argument("parser", help="importable parser (or function returning parser)")
     parser.add_argument("--version", action="version", version="%(prog)s " + __version__)
     parser.add_argument("-s", "--shell", default=SUPPORTED_SHELLS[0], choices=SUPPORTED_SHELLS)
-    parser.add_argument("-o", "--output", help="write output to file instead of stdout", type=Path)
+    parser.add_argument("-o", "--output", default='-', help="output file (- for stdout)",
+                        type=Path)
     parser.add_argument("--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`)")
@@ -67,7 +56,16 @@ def main(argv=None):
         other_parser = other_parser()
     if args.prog:
         other_parser.prog = args.prog
-    with extract_stdout(args.output) as stdout:
+
+    @contextmanager
+    def _open(out_path):
+        if str(out_path) in ("-", "stdout"):
+            yield sys.stdout
+        else:
+            with out_path.open('w') as fd:
+                yield fd
+
+    with _open(args.output) as fd:
         print(
             complete(other_parser, shell=args.shell, root_prefix=args.prefix
-                     or args.parser.split(".", 1)[0], preamble=args.preamble), file=stdout)
+                     or args.parser.split(".", 1)[0], preamble=args.preamble), file=fd)

tests/test_shtab.py

@@ -9,7 +9,7 @@
 import pytest
 
 import shtab
-from shtab.main import extract_stdout, get_main_parser, main
+from shtab.main import get_main_parser, main
 
 fix_shell = pytest.mark.parametrize("shell", shtab.SUPPORTED_SHELLS)
 
@@ -82,6 +82,31 @@ def test_main_self_completion(shell, caplog, capsys):
     assert not caplog.record_tuples
 
 
+@pytest.mark.parametrize('output', ["-", "stdout", "test.txt"])
+@fix_shell
+def test_main_output_path(shell, caplog, capsys, change_dir, output):
+    assert not capsys.readouterr().out
+    with caplog.at_level(logging.INFO):
+        try:
+            main(["-s", shell, "shtab.main.get_main_parser", "-o", output])
+        except SystemExit:
+            pass
+
+    captured = capsys.readouterr()
+    assert not captured.err
+    expected = {
+        "bash": "complete -o filenames -F _shtab_shtab shtab", "zsh": "_shtab_shtab_commands()",
+        "tcsh": "complete shtab"}
+
+    if output in ("-", "stdout"):
+        assert expected[shell] in captured.out
+    else:
+        assert not captured.out
+        assert expected[shell] in (change_dir / output).read_text()
+
+    assert not caplog.record_tuples
+
+
 @fix_shell
 def test_prog_override(shell, caplog, capsys):
     with caplog.at_level(logging.INFO):
@@ -342,17 +367,3 @@ def test_path_completion_after_redirection(caplog, change_dir):
         shell.test('"${COMPREPLY[@]}" = "test_file.txt"', f"Redirection {redirection} failed")
 
     assert not caplog.record_tuples
-
-
-def test_extract_stdout(tmp_path):
-    path = tmp_path / "completions"
-    with extract_stdout(path) as output:
-        output.write("completion")
-    assert path.read_text() == "completion"
-
-
-def test_extract_stdout_empty(capsys):
-    with extract_stdout(None) as output:
-        output.write("completion")
-    captured = capsys.readouterr()
-    assert captured.out == "completion"