ast-grep automatic language selection (#1394)

* prep ts docs

* update gitignore rules

* automatic language selection

* document mapping

* typo

* revert api change

* typo

* fix sample

* no model needed

* fix test

Author: pelikhan

.gitignore

@@ -39,4 +39,5 @@ packages/sample/src/xpai/*.json
 packages/sample/src/xpai/*.csv
 packages/sample/src/mlads/*
 packages/core/src/dbg.ts
-
+TypeScript/
+react/

docs/astro.config.mjs

@@ -58,6 +58,9 @@ is at https://microsoft.github.io/genaiscript/reference/scripts.md
 - save generated code in the "./genaisrc" folder with ".genai.mts" extension
 `,
                     pageSeparator: "\n\n=|=|=|=|=|=\n\n",
+                    minify: {
+                        customSelectors: ["picture"]
+                    },
                     promote: ["index*", "getting-started*", "!*/*"],
                     customSets: [
                         {

docs/src/content/docs/reference/scripts/ast-grep.mdx

@@ -132,9 +132,9 @@ This version of `ast-grep` [supports the following built-in languages](https://a
 
 - Html
 - JavaScript
+- TypeScript
 - Tsx
 - Css
-- TypeScript
 
 The following languages require installing an additional package:
 
@@ -144,6 +144,9 @@ The following languages require installing an additional package:
 npm install -D @ast-grep/lang-c
 ```
 
+- C++, `@ast-grep/lang-cpp`
+- Python, `@ast-grep/lang-python`
+- C#, `@ast-grep/lang-csharp`
 - SQL, `@ast-grep/lang-sql`
 - Angular, `@ast-grep/lang-angular`
 
@@ -153,6 +156,30 @@ If your language is not supported, go to [ast-grep langs](https://github.com/ast
 
 :::
 
+### Filename extension mapping
+
+The following file extensions are mapped to the corresponding languages:
+
+- HTML: `html`, `htm`
+- JavaScript: `cjs`, `mjs`, `js`
+- TypeScript: `cts`, `mts`, `ts`
+- TSX: `tsx`
+- CSS: `css`    
+- c: `c`
+- cpp: `cpp`, `cxx`, `h`, `hpp`, `hxx`
+- python: `py`
+- C#: `cs`
+- sql: `sql`
+
+### Overriding the language selection
+
+GenAIScript has default mappings from well-known file extensions to languages.
+However, you can override this by passing the `lang` option to the `search` method.
+
+```ts "{ lang: "ts" }"
+const { matches } = await sg.search("ts", "src/fib.ts", {...}, { lang: "ts" })
+```
+
 ## Learning ast-grep
 
 There is a learning curve to grasp the query language of `ast-grep`.

docs/src/content/docs/reference/scripts/diagrams.md

@@ -62,3 +62,9 @@ graph LR
 ````
 
 and it gets rendered automatically once you install the extension.
+
+```mermaid
+graph LR
+  A[Master] --> C[New Commit]
+  B[Feature Branch] --> C
+```

genaisrc/docs.genai.mts

@@ -34,13 +34,17 @@ script({
             default: true,
             description: "Update existing docs.",
         },
+        maxFiles: {
+            type: "integer",
+            description: "Maximum number of files to process.",
+        },
     },
 })
 const { output, dbg, vars } = env
 let { files } = env
-const { applyEdits, diff, pretty, missing, update } = vars
+const { applyEdits, diff, pretty, missing, update, maxFiles } = vars
 
-dbg({ applyEdits, diff, pretty, missing, update })
+dbg({ applyEdits, diff, pretty, missing, update, maxFiles })
 
 if (!missing && !update) cancel(`not generating or updating docs, exiting...`)
 
@@ -60,6 +64,14 @@ if (diffFiles?.length) {
     )
     dbg(`diff filtered files: ${files.length}`)
 }
+
+if (maxFiles && files.length > maxFiles) {
+    dbg(`random slicing files to ${maxFiles}`)
+    files = parsers.tidyData(files, {
+        sliceSample: maxFiles,
+    }) as WorkspaceFile[]
+}
+
 const sg = await host.astGrep()
 const stats = []
 for (const file of files) {
@@ -123,7 +135,7 @@ async function generateDocs(file: WorkspaceFile, fileStats: any) {
                 },
             },
         },
-        { diff: gitDiff }
+        { diff: gitDiff, applyGitIgnore: false }
     )
     dbg(`found ${missingDocs.length} missing docs`)
     const edits = sg.changeset()
@@ -206,19 +218,16 @@ async function updateDocs(file: WorkspaceFile, fileStats: any) {
     const { matches } = await sg.search(
         "ts",
         file.filename,
-        {
-            rule: {
-                kind: "export_statement",
-                follows: {
-                    kind: "comment",
-                    stopBy: "neighbor",
-                },
-                has: {
-                    kind: "function_declaration",
-                },
-            },
-        },
-        { diff: gitDiff }
+        YAML`
+rule: 
+  kind: "export_statement"
+  follows: 
+    kind: "comment"
+    stopBy: neighbor
+  has:
+      kind: "function_declaration"
+`,
+        { diff: gitDiff, applyGitIgnore: false }
     )
     dbg(`found ${matches.length} docs to update`)
     const edits = sg.changeset()

genaisrc/update-tree-sitter-queries.genai.mts

@@ -1,4 +1,5 @@
 import { Octokit } from "octokit"
+script({ model: "none" })
 
 const files: Record<string, string> = {}
 const downloadScm = async (repo: string, name: string) => {

packages/core/src/astgrep.test.ts

@@ -2,6 +2,7 @@ import { beforeEach, describe, test } from "node:test"
 import assert from "node:assert/strict"
 import { astGrepFindFiles, astGrepParse } from "./astgrep"
 import { TestHost } from "./testhost"
+import { dedent } from "./indent"
 
 describe("astgrep", () => {
     beforeEach(() => {
@@ -35,4 +36,44 @@ describe("astgrep", () => {
         const result = await astGrepParse(file, { lang: "js" })
         assert.equal(result, undefined)
     })
+
+    test("parse C++ file", async () => {
+        const file: WorkspaceFile = {
+            filename: "test.cpp",
+            content: dedent`
+            #include <iostream>
+            
+            int main() {
+                std::cout << 'Hello, world!' << std::endl;
+                return 0;
+            }
+            `,
+        }
+        const result = await astGrepParse(file)
+        assert(result)
+    })
+    test("parse TypeScript file", async () => {
+        const file: WorkspaceFile = {
+            filename: "test.ts",
+            content: "const x: number = 1;",
+        }
+        const result = await astGrepParse(file)
+        assert(result)
+    })
+    test("parse python file", async () => {
+        const file: WorkspaceFile = {
+            filename: "test.py",
+            content: "x = 1",
+        }
+        const result = await astGrepParse(file)
+        assert(result)
+    })
+    test("parse C file", async () => {
+        const file: WorkspaceFile = {
+            filename: "test.c",
+            content: "#include <stdio.h>",
+        }
+        const result = await astGrepParse(file)
+        assert(result)
+    })
 })

packages/core/src/astgrep.ts

@@ -1,5 +1,6 @@
 import debug from "debug"
 const dbg = debug("genaiscript:astgrep")
+const dbgLang = debug("genaiscript:astgrep:lang")
 
 import { CancellationOptions, checkCancelled } from "./cancellation"
 import { CancelError, errorMessage } from "./error"
@@ -94,12 +95,10 @@ export async function astGrepFindFiles(
     }
     const diffFiles = diffResolve(diff)
 
-    dbg(`finding files with ${lang} %O`, matcher)
+    dbg(`search %O`, matcher)
     if (diffFiles?.length) dbg(`diff files: ${diffFiles.length}`)
     const { findInFiles } = await import("@ast-grep/napi")
     checkCancelled(cancellationToken)
-    const sglang = await resolveLang(lang)
-    dbg(`resolving language: ${lang}`)
 
     let paths = await host.findFiles(glob, options)
     if (!paths?.length) {
@@ -129,6 +128,7 @@ export async function astGrepFindFiles(
     const p = new Promise<number>(async (resolve, reject) => {
         let i = 0
         let n: number = undefined
+        const sglang = await resolveLang(lang)
         n = await findInFiles(
             sglang,
             {
@@ -230,7 +230,7 @@ export async function astGrepWriteRootEdits(
  */
 export async function astGrepParse(
     file: WorkspaceFile,
-    options?: { lang?: SgLang } & CancellationOptions
+    options?: { lang?: SgLang | Record<string, SgLang> } & CancellationOptions
 ): Promise<SgRoot> {
     const { cancellationToken } = options || {}
     if (file.encoding) {
@@ -249,7 +249,6 @@ export async function astGrepParse(
     dbg(`parsing file: ${filename}`)
     const { parseAsync } = await import("@ast-grep/napi")
     const lang = await resolveLang(options?.lang, filename)
-    dbg(`resolving language for file: ${filename}`)
     if (!lang) {
         return undefined
     }
@@ -259,65 +258,91 @@ export async function astGrepParse(
     return root
 }
 
-async function resolveLang(lang: SgLang, filename?: string) {
+async function resolveLang(
+    lang: SgLang | Record<string, SgLang>,
+    filename?: string
+) {
     const { Lang } = await import("@ast-grep/napi")
-    if (lang === "html") {
-        return Lang.Html
-    }
-    if (lang === "js") {
-        return Lang.JavaScript
-    }
-    if (lang === "ts") {
-        return Lang.TypeScript
+
+    const norm = (l: string) => l.toLowerCase().replace(/^\./, "")
+
+    // pre-compiled with ast-grep
+    const builtins: any = {
+        html: Lang.Html,
+        htm: Lang.Html,
+        cjs: Lang.JavaScript,
+        mjs: Lang.JavaScript,
+        js: Lang.JavaScript,
+        cts: Lang.TypeScript,
+        mts: Lang.TypeScript,
+        ts: Lang.TypeScript,
+        tsx: Lang.Tsx,
+        css: Lang.Css,
     }
-    if (lang === "tsx") {
-        return Lang.Tsx
+
+    const dynamics: any = {
+        h: "c",
+        c: "c",
+        cpp: "cpp",
+        hpp: "cpp",
+        hxx: "cpp",
+        cxx: "cpp",
+        cs: "csharp",
+        py: "python",
+        sql: "sql",
     }
-    if (lang === "css") {
-        return Lang.Css
+
+    const forbidden = ["bin", "exe", "dll"]
+
+    // user provided a string
+    if (typeof lang === "string") {
+        lang = norm(lang)
+        dbgLang(`resolving language ${lang}`)
+        const builtin = builtins[lang]
+        if (builtin) return builtin
+        else return await loadDynamicLanguage(lang)
     }
-    if (lang) {
-        return await loadDynamicLanguage(lang.toLowerCase())
+
+    if (!filename) {
+        dbgLang(`filename not provided`)
+        throw new Error("filename is required to resolve language")
     }
 
     if (filename) {
-        dbg(`resolving language based on filename: ${filename}`)
-        if (/\.m?js$/i.test(filename)) {
-            return Lang.JavaScript
-        }
-        if (/\.m?ts$/i.test(filename)) {
-            return Lang.TypeScript
-        }
-        if (/\.(j|t)sx$/i.test(filename)) {
-            return Lang.Tsx
-        }
-        if (/\.html$/i.test(filename)) {
-            return Lang.Html
-        }
-        if (/\.css$/i.test(filename)) {
-            return Lang.Css
-        }
-        return await loadDynamicLanguage(
-            extname(filename).slice(1).toLowerCase()
-        )
+        const ext = norm(extname(filename))
+        dbgLang(`resolving language for ${ext}`)
+
+        // known builtins
+        const builtin = builtins[ext]
+        if (builtin) return builtin
+
+        // known dynamics
+        const dynamic = dynamics[ext]
+        if (dynamic) return await loadDynamicLanguage(dynamic)
+
+        if (forbidden.includes(ext)) return undefined
+
+        // try our luck
+        return await loadDynamicLanguage(ext)
     }
 
+    dbgLang(`language not resolved`, { lang, filename })
     throw new Error("language not resolved")
 }
 
 const loadedDynamicLanguages = new Set<string>()
 async function loadDynamicLanguage(langName: string) {
     if (!loadedDynamicLanguages.has(langName)) {
-        dbg(`loading language: ${langName}`)
+        dbgLang(`loading language: ${langName}`)
         const { registerDynamicLanguage } = await import("@ast-grep/napi")
         try {
             const dynamicLang = (await import(`@ast-grep/lang-${langName}`))
                 .default
             registerDynamicLanguage({ [langName]: dynamicLang })
             loadedDynamicLanguages.add(langName)
-            dbg(`language ${langName} registered `)
+            dbgLang(`language ${langName} registered `)
         } catch (err) {
-            dbg(`error loading language ${langName}: ${errorMessage(err)}`)
+            dbgLang(`error loading language ${langName}: ${errorMessage(err)}`)
             throw Error(
                 `@ast-grep/lang-${langName} package failed to load, please install it using 'npm install -D @ast-grep/lang-${langName}'`
             )