Add ability to write completions to the file

docs/use.md

@@ -74,7 +74,7 @@ Below are various examples of enabling `shtab`'s own tab completion scripts.
     ```sh
     mkdir -p ~/.zsh/completions
     fpath=($fpath ~/.zsh/completions)  # must be before `compinit` lines
-    shtab --shell=zsh shtab.main.get_main_parser > ~/.zsh/completions/_shtab
+    shtab --shell=zsh shtab.main.get_main_parser -o ~/.zsh/completions/_shtab
     ```
 
 === "tcsh"

shtab/main.py

@@ -1,19 +1,34 @@
 import argparse
+import contextlib
 import logging
 import os
 import sys
 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("--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`)")
@@ -52,6 +67,7 @@ def main(argv=None):
         other_parser = other_parser()
     if args.prog:
         other_parser.prog = args.prog
-    print(
-        complete(other_parser, shell=args.shell, root_prefix=args.prefix
-                 or args.parser.split(".", 1)[0], preamble=args.preamble))
+    with extract_stdout(args.output) as stdout:
+        print(
+            complete(other_parser, shell=args.shell, root_prefix=args.prefix
+                     or args.parser.split(".", 1)[0], preamble=args.preamble), file=stdout)

tests/test_shtab.py

@@ -9,7 +9,7 @@
 import pytest
 
 import shtab
-from shtab.main import get_main_parser, main
+from shtab.main import extract_stdout, get_main_parser, main
 
 fix_shell = pytest.mark.parametrize("shell", shtab.SUPPORTED_SHELLS)
 
@@ -342,3 +342,17 @@ 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"