nixpkgs/pkgs/applications/editors/vim/plugins/utils/nvim-treesitter/update.py

#!/usr/bin/env nix-shell
#!nix-shell ./update-shell.nix -i python

import json
import logging
import os
import subprocess
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import requests

log = logging.getLogger("vim-updater")

NURR_JSON_URL = (
    "https://raw.githubusercontent.com/nvim-neorocks/nurr/main/tree-sitter-parsers.json"
)


def generate_grammar(lang, parser_info):
    """Generate grammar for a language based on the parser info"""
    if "install_info" not in parser_info:
        log.warning(f"Parser {lang} does not have install_info, skipping")
        return ""

    install_info = parser_info["install_info"]

    url = install_info["url"]
    rev = install_info["revision"]

    generated = f"""  {lang} = buildGrammar {{
    language = "{lang}";
    version = "0.0.0+rev={rev[:7]}";
    src = """

    generated += subprocess.check_output(
        ["nurl", url, rev, "--indent=4"], text=True
    )
    generated += ";"

    location = install_info.get("location", "")
    if location:
        generated += f"""
    location = "{location}";"""

    if install_info.get("generate", False):
        generated += """
    generate = true;"""

    generated += f"""
    meta.homepage = "{url}";
  }};
"""

    return generated


def fetch_nurr_parsers():
    """Fetch the parser information from nurr repository"""
    log.info("Fetching parser data from %s", NURR_JSON_URL)

    headers = {}
    github_token = os.environ.get("GITHUB_TOKEN")
    if github_token:
        log.info("Using GITHUB_TOKEN for authentication")
        headers["Authorization"] = f"token {github_token}"
    else:
        log.warning("No GITHUB_TOKEN found. GitHub API requests may be rate-limited.")

    response = requests.get(NURR_JSON_URL, headers=headers, timeout=30)
    response.raise_for_status()
    data = response.json()

    try:
        parsers = data["parsers"]
    except KeyError:
        raise ValueError(
            "Unexpected response from NURR:\n" + json.dumps(data, indent=2)
        )
    log.info(f"Successfully fetched {len(parsers)} parsers")
    return parsers


def process_parser_info(parser_info):
    """Process a single parser info entry and generate grammar for it"""
    return generate_grammar(parser_info["lang"], parser_info)


def update_grammars():
    """Update grammar definitions using nurr's parser information"""
    parsers_info = fetch_nurr_parsers()

    generated_file = """# generated by pkgs/applications/editors/vim/plugins/utils/nvim-treesitter/update.py
# Using parser data from https://github.com/nvim-neorocks/nurr/blob/main/tree-sitter-parsers.json

{
  buildGrammar,
  """

    nurl_output = subprocess.check_output(["nurl", "-Ls", ","], text=True).strip()
    indented_output = nurl_output.replace(",", ",\n  ")
    generated_file += indented_output
    generated_file += """,
}:

{
"""

    # Process parsers in parallel for better performance
    with ThreadPoolExecutor(max_workers=5) as executor:
        for generated in executor.map(process_parser_info, parsers_info):
            generated_file += generated

    generated_file += "}\n"
    return generated_file


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")

    generated = update_grammars()
    output_path = Path(__file__).parent.parent.parent / "nvim-treesitter/generated.nix"
    log.info("Writing output to %s", output_path)
    with open(output_path, "w") as f:
        f.write(generated)
    log.info("Successfully updated grammar definitions")