1 files changed, 53 insertions, 32 deletions
diff --git a/pkgs/applications/editors/vim/plugins/nvim-treesitter/update.py b/pkgs/applications/editors/vim/plugins/nvim-treesitter/update.py
index bbacc6959c2..4b70d8555dd 100755
--- a/pkgs/applications/editors/vim/plugins/nvim-treesitter/update.py
+++ b/pkgs/applications/editors/vim/plugins/nvim-treesitter/update.py
@@ -2,36 +2,20 @@
 #!nix-shell update-shell.nix -i python
 
 import json
+import logging
 import subprocess
 from concurrent.futures import ThreadPoolExecutor
-from os import environ
-from os.path import dirname, join
-
-lockfile = json.load(open(join(environ["NVIM_TREESITTER"], "lockfile.json")))
-
-configs = json.loads(
-    subprocess.check_output(
-        [
-            "nvim",
-            "--headless",
-            "-u",
-            "NONE",
-            "+lua io.write(vim.json.encode(require('nvim-treesitter.parsers').get_parser_configs()))",
-            "+quit!",
-        ]
-    )
-)
-
-
-def generate_grammar(item):
-    lang, lock = item
-    cfg = configs.get(lang)
-    if not cfg:
-        return ""
+import os
+import sys
+from os.path import join
 
+log = logging.getLogger("vim-updater")
+
+
+def generate_grammar(lang, rev, cfg):
+    """Generate grammar for a language"""
     info = cfg["install_info"]
     url = info["url"]
-    rev = lock["revision"]
 
     generated = f"""  {lang} = buildGrammar {{
     language = "{lang}";
@@ -58,20 +42,57 @@ def generate_grammar(item):
     return generated
 
 
-generated_file = """# generated by pkgs/applications/editors/vim/plugins/nvim-treesitter/update.py
+def update_grammars(nvim_treesitter_dir: str):
+    """
+    The lockfile contains just revisions so we start neovim to dump the
+    grammar information in a better format
+    """
+    # the lockfile
+    cmd = [
+        "nvim",
+        "--headless",
+        "-u",
+        "NONE",
+        "--cmd",
+        f"set rtp^={nvim_treesitter_dir}",
+        "+lua io.write(vim.json.encode(require('nvim-treesitter.parsers').get_parser_configs()))",
+        "+quit!",
+    ]
+    log.debug("Running command: %s", ' '.join(cmd))
+    configs = json.loads(subprocess.check_output(cmd))
+
+    generated_file = """# generated by pkgs/applications/editors/vim/plugins/nvim-treesitter/update.py
 
 { buildGrammar, """
 
-generated_file += subprocess.check_output(["nurl", "-Ls", ", "], text=True)
+    generated_file += subprocess.check_output(["nurl", "-Ls", ", "], text=True)
 
-generated_file += """ }:
+    generated_file += """ }:
 
 {
 """
 
-for generated in ThreadPoolExecutor().map(generate_grammar, lockfile.items()):
-    generated_file += generated
+    lockfile_path = os.path.join(nvim_treesitter_dir, "lockfile.json")
+    log.debug("Opening %s", lockfile_path)
+    with open(lockfile_path) as lockfile_fd:
+        lockfile = json.load(lockfile_fd)
+
+        def _generate_grammar(item):
+            lang, lock = item
+            cfg = configs.get(lang)
+            if not cfg:
+                return ""
+            return generate_grammar(lang, lock["revision"], cfg)
+
+        for generated in ThreadPoolExecutor(max_workers=5).map(
+            _generate_grammar, lockfile.items()
+        ):
+            generated_file += generated
+
+        generated_file += "}\n"
+    return generated_file
 
-generated_file += "}\n"
 
-open(join(dirname(__file__), "generated.nix"), "w").write(generated_file)
+if __name__ == "__main__":
+    generated = update_grammars(sys.args[1])
+    open(join(os.path.dirname(__file__), "generated.nix"), "w").write(generated)