vimPluginsUpdater: fix treesitter updates

4 files changed, 72 insertions, 55 deletions

diff --git a/maintainers/scripts/pluginupdate.py b/maintainers/scripts/pluginupdate.py
index 18b354f0be9..5ceaab8db90 100644
--- a/maintainers/scripts/pluginupdate.py
+++ b/maintainers/scripts/pluginupdate.py
@@ -327,7 +327,6 @@ def run_nix_expr(expr, nixpkgs: str):
     :param expr nix expression to fetch current plugins
     :param nixpkgs Path towards a nixpkgs checkout
     '''
-    # local_pkgs = str(Path(__file__).parent.parent.parent)
     with CleanEnvironment(nixpkgs) as nix_path:
         cmd = [
             "nix",
@@ -341,8 +340,8 @@ def run_nix_expr(expr, nixpkgs: str):
             "--nix-path",
             nix_path,
         ]
-        log.debug("Running command %s", " ".join(cmd))
-        out = subprocess.check_output(cmd)
+        log.debug("Running command: %s", " ".join(cmd))
+        out = subprocess.check_output(cmd, timeout=90)
         data = json.loads(out)
         return data
 
@@ -572,7 +571,6 @@ class CleanEnvironment(object):
         self.empty_config = NamedTemporaryFile()
         self.empty_config.write(b"{}")
         self.empty_config.flush()
-        # os.environ["NIXPKGS_CONFIG"] = self.empty_config.name
         return f"localpkgs={self.local_pkgs}"
 
     def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
diff --git a/pkgs/applications/editors/vim/plugins/nvim-treesitter/update.py b/pkgs/applications/editors/vim/plugins/nvim-treesitter/update.py
index 37414b07353..6d9fd4de098 100755
--- a/pkgs/applications/editors/vim/plugins/nvim-treesitter/update.py
+++ b/pkgs/applications/editors/vim/plugins/nvim-treesitter/update.py
@@ -2,34 +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
-
-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}";
@@ -56,7 +42,24 @@ def generate_grammar(item):
     return generated
 
 
-def update_grammars(lockfile: str):
+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
 
@@ -68,14 +71,27 @@ def update_grammars(lockfile: str):
 
     {
     """
-    for generated in ThreadPoolExecutor().map(generate_grammar, lockfile.items()):
-        generated_file += generated
-        generated_file += "}\n"
-        generated_file += "}\n"
 
-    open(join(dirname(__file__), "generated.nix"), "w").write(generated_file)
+    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
 
 
 if __name__ == "__main__":
-    # TODO add lockfile
-    update_grammars()
+    generated = update_grammars(sys.args[1])
+    open(join(os.path.dirname(__file__), "generated.nix"), "w").write(generated)
diff --git a/pkgs/applications/editors/vim/plugins/update.py b/pkgs/applications/editors/vim/plugins/update.py
index 438a231ec47..b658ca7e2bd 100755
--- a/pkgs/applications/editors/vim/plugins/update.py
+++ b/pkgs/applications/editors/vim/plugins/update.py
@@ -23,11 +23,12 @@ import os
 import logging
 import textwrap
 import json
+import subprocess
 from typing import List, Tuple
 from pathlib import Path
 
 
-log = logging.getLogger()
+log = logging.getLogger("vim-updater")
 
 sh = logging.StreamHandler()
 formatter = logging.Formatter("%(name)s:%(levelname)s: %(message)s")
@@ -39,15 +40,14 @@ ROOT = Path(os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe
 import pluginupdate
 import importlib
 from pluginupdate import run_nix_expr, PluginDesc
-from treesitter import update_grammars
+import treesitter
 
 
 HEADER = (
     "# GENERATED by ./pkgs/applications/editors/vim/plugins/update.py. Do not edit!"
 )
 
-NIXPKGS_NVIMTREESITTER_FOLDER = \
-    "pkgs/applications/editors/vim/plugins/nvim-treesitter/generated.nix"
+NIXPKGS_NVIMTREESITTER_FOLDER = "pkgs/applications/editors/vim/plugins/nvim-treesitter"
 
 
 class VimEditor(pluginupdate.Editor):
@@ -58,8 +58,7 @@ class VimEditor(pluginupdate.Editor):
     ):
         sorted_plugins = sorted(plugins, key=lambda v: v[0].name.lower())
         nvim_treesitter_rev = pluginupdate.run_nix_expr(
-            "(import <localpkgs> { }).vimPlugins.nvim-treesitter.src.rev",
-            self.nixpkgs
+            "(import <localpkgs> { }).vimPlugins.nvim-treesitter.src.rev", self.nixpkgs
         )
 
         with open(outfile, "w+") as f:
@@ -78,7 +77,8 @@ class VimEditor(pluginupdate.Editor):
                 content = self.plugin2nix(pdesc, plugin)
                 f.write(content)
                 if (
-                    plugin.name == "nvim-treesitter" and plugin.commit != nvim_treesitter_rev
+                    plugin.name == "nvim-treesitter"
+                    and plugin.commit != nvim_treesitter_rev
                 ):
                     self.nvim_treesitter_updated = True
             f.write("\n}\n")
@@ -126,13 +126,19 @@ class VimEditor(pluginupdate.Editor):
     def update(self, args):
         pluginupdate.update_plugins(self, args)
 
+        # TODO this should probably be skipped when running outside a nixpkgs checkout
         if self.nvim_treesitter_updated:
             print("updating nvim-treesitter grammars")
-            nvim_treesitter_dir = ROOT.joinpath("nvim-treesitter")
-            lockfile = os.path.join(args.nixpkgs.join(NIXPKGS_NVIMTREESITTER_FOLDER, "lockfile.json"))
-            lockfile = json.load(open(lockfile))
+            cmd = [
+                "nix", "build",
+                "vimPlugins.nvim-treesitter.src", "-f", self.nixpkgs
+                , "--print-out-paths"
+            ]
+            log.debug("Running command: %s", " ".join(cmd))
+            nvim_treesitter_dir = subprocess.check_output(cmd, text=True, timeout=90).strip()
 
-            nvim_treesitter.update_grammars(lockfile)
+            generated = treesitter.update_grammars(nvim_treesitter_dir)
+            open(os.path.join(args.nixpkgs, "generated.nix"), "w").write(generated)
 
             if self.nixpkgs_repo:
                 index = self.nixpkgs_repo.index
@@ -147,13 +153,14 @@ class VimEditor(pluginupdate.Editor):
 
 
 def main():
-
     global luaPlugins
 
     log.debug(f"Loading from {ROOT}/../get-plugins.nix")
     with open(f"{ROOT}/../get-plugins.nix") as f:
         GET_PLUGINS = f.read()
-    editor = VimEditor("vim", Path("pkgs/applications/editors/vim/plugins"), GET_PLUGINS)
+    editor = VimEditor(
+        "vim", Path("pkgs/applications/editors/vim/plugins"), GET_PLUGINS
+    )
     editor.run()
 
 
diff --git a/pkgs/applications/editors/vim/plugins/updater.nix b/pkgs/applications/editors/vim/plugins/updater.nix
index d17247b3f8a..afa245d9be9 100644
--- a/pkgs/applications/editors/vim/plugins/updater.nix
+++ b/pkgs/applications/editors/vim/plugins/updater.nix
@@ -4,17 +4,12 @@
 , python3Packages
 , lib
 , nix-prefetch-git
+, nurl
 
 # optional
 , vimPlugins
 , neovim
 }:
-let
-  my_neovim = neovim.override {
-    configure.packages.all.start = [ vimPlugins.nvim-treesitter ];
-  };
-
-in
 buildPythonApplication {
   format = "other";
   pname = "vim-plugins-updater";
@@ -39,7 +34,8 @@ buildPythonApplication {
     cp ${../../../../../maintainers/scripts/pluginupdate.py} $out/lib/pluginupdate.py
 
     # wrap python scripts
-    makeWrapperArgs+=( --prefix PATH : "${lib.makeBinPath [ nix nix-prefetch-git my_neovim ]}" --prefix PYTHONPATH : "$out/lib" )
+    makeWrapperArgs+=( --prefix PATH : "${lib.makeBinPath [
+      nix nix-prefetch-git neovim nurl ]}" --prefix PYTHONPATH : "$out/lib" )
     wrapPythonPrograms
   '';