Share option gathering between Bash and zsh completions

devscripts/bash-completion.py

@@ -2,60 +2,31 @@
 
 # Allow direct execution
 import os
-import shlex
 import sys
-from dataclasses import dataclass
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 
 import yt_dlp
+from devscripts.utils import gather_completion_opts
 
 BASH_COMPLETION_FILE = 'completions/bash/yt-dlp'
 BASH_COMPLETION_TEMPLATE = 'devscripts/bash-completion.in'
 
 
-@dataclass
-class YtDlpOpts:
-    opts: tuple[str, ...]
-    file_opts: tuple[str, ...]
-    dir_opts: tuple[str, ...]
+def build_completion(opt_parser):
+    opts = gather_completion_opts(opt_parser)
 
+    with open(BASH_COMPLETION_TEMPLATE) as f:
+        template = f.read()
 
-def yt_dlp_options():
-    opts, long_opts = [], []
-    file_opts, dir_opts = [], []
-
-    for opt in yt_dlp.parseOpts(ignore_config_files=True)[0]._get_all_options():
-        opts.extend(opt._short_opts)
-        long_opts.extend(opt._long_opts)
-        if opt.metavar in {'FILE', 'PATH', 'CERTFILE', 'KEYFILE'}:
-            file_opts.extend(opt._short_opts)
-            file_opts.extend(opt._long_opts)
-        elif opt.metavar == 'DIR':
-            dir_opts.extend(opt._short_opts)
-            dir_opts.extend(opt._long_opts)
-
-    opts.sort()
-    long_opts.sort()
-    opts.extend(long_opts)
+    template = template.replace('YT_DLP_FILE_OPTS_CASE', ' | '.join(opts.file_opts))
+    template = template.replace('YT_DLP_DIR_OPTS_CASE', ' | '.join(opts.dir_opts))
+    template = template.replace('YT_DLP_OPTS_ARRAY', ' '.join(opts.opts))
 
-    return YtDlpOpts(tuple(opts), tuple(file_opts), tuple(dir_opts))
+    with open(BASH_COMPLETION_FILE, 'w') as f:
+        f.write(template)
 
 
-def build_completion():
-    with open(BASH_COMPLETION_TEMPLATE) as f:
-        template = f.read()
-    with open(BASH_COMPLETION_FILE, 'w') as f:
-        # just using the special char
-        opts = yt_dlp_options()
-        f.write(template.replace(
-            'YT_DLP_FILE_OPTS_CASE', ' | '.join(shlex.quote(o) for o in opts.file_opts),
-        ).replace(
-            'YT_DLP_DIR_OPTS_CASE', ' | '.join(shlex.quote(o) for o in opts.dir_opts),
-        ).replace(
-            'YT_DLP_OPTS_ARRAY', ' '.join(shlex.quote(o) for o in opts.opts),
-        ))
-
-
-build_completion()
+parser = yt_dlp.parseOpts(ignore_config_files=True)[0]
+build_completion(parser)

devscripts/utils.py

@@ -2,7 +2,9 @@ import argparse
 import datetime as dt
 import functools
 import re
+import shlex
 import subprocess
+from dataclasses import dataclass
 
 
 def read_file(fname):
@@ -64,3 +66,36 @@ def run_process(*args, **kwargs):
         kwargs.setdefault('encoding', 'utf-8')
         kwargs.setdefault('errors', 'replace')
     return subprocess.run(args, **kwargs)
+
+
+@dataclass
+class YtDlpCompletionOpts:
+    opts: tuple[str, ...]
+    file_opts: tuple[str, ...]
+    dir_opts: tuple[str, ...]
+
+
+def gather_completion_opts(yt_dlp_opt_parser):
+    opts, long_opts = [], []
+    file_opts, dir_opts = [], []
+
+    for opt in yt_dlp_opt_parser._get_all_options():
+        opts.extend(opt._short_opts)
+        long_opts.extend(opt._long_opts)
+        if opt.metavar in {'FILE', 'PATH', 'CERTFILE', 'KEYFILE'}:
+            file_opts.extend(opt._short_opts)
+            file_opts.extend(opt._long_opts)
+        elif opt.metavar == 'DIR':
+            dir_opts.extend(opt._short_opts)
+            dir_opts.extend(opt._long_opts)
+
+    opts.sort()
+    long_opts.sort()
+    opts.extend(long_opts)
+
+    # Escape after sorting.
+    return YtDlpCompletionOpts(
+        tuple(shlex.quote(o) for o in opts),
+        tuple(shlex.quote(o) for o in file_opts),
+        tuple(shlex.quote(o) for o in dir_opts),
+    )

devscripts/zsh-completion.py

@@ -8,46 +8,21 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 
 import yt_dlp
+from devscripts.utils import gather_completion_opts
 
 ZSH_COMPLETION_FILE = 'completions/zsh/_yt-dlp'
 ZSH_COMPLETION_TEMPLATE = 'devscripts/zsh-completion.in'
 
 
 def build_completion(opt_parser):
-    opts = [opt for group in opt_parser.option_groups
-            for opt in group.option_list]
-    opts_file = [opt for opt in opts if opt.metavar == 'FILE']
-    opts_dir = [opt for opt in opts if opt.metavar == 'DIR']
-    opts_path = [opt for opt in opts if opt.metavar == 'PATH']
-
-    fileopts = []
-    for opt in opts_file:
-        if opt._short_opts:
-            fileopts.extend(opt._short_opts)
-        if opt._long_opts:
-            fileopts.extend(opt._long_opts)
-
-    for opt in opts_path:
-        if opt._short_opts:
-            fileopts.extend(opt._short_opts)
-        if opt._long_opts:
-            fileopts.extend(opt._long_opts)
-
-    diropts = []
-    for opt in opts_dir:
-        if opt._short_opts:
-            diropts.extend(opt._short_opts)
-        if opt._long_opts:
-            diropts.extend(opt._long_opts)
-
-    flags = [opt.get_opt_string() for opt in opts]
+    opts = gather_completion_opts(opt_parser)
 
     with open(ZSH_COMPLETION_TEMPLATE) as f:
         template = f.read()
 
-    template = template.replace('{{fileopts}}', '|'.join(fileopts))
-    template = template.replace('{{diropts}}', '|'.join(diropts))
-    template = template.replace('{{flags}}', ' '.join(flags))
+    template = template.replace('{{fileopts}}', '|'.join(opts.file_opts))
+    template = template.replace('{{diropts}}', '|'.join(opts.dir_opts))
+    template = template.replace('{{flags}}', ' '.join(opts.opts))
 
     with open(ZSH_COMPLETION_FILE, 'w') as f:
         f.write(template)