Skip to content

Commit e875555

Browse files
authored
Ast dbg (#1373)
* more ast-grep fixes * ✨ feat: enhance file diffing and add cowsay functionality Added string handling to diff creation. Introduced a cowsay module. * ✨ feat: enhance function docstring generation logic Improve docstring generation with concise syntax and error handling. * updated script * ✨ Enhance script with doc generation and edits Added function doc generation, edit commitment, and normalization.
1 parent fce5808 commit e875555

8 files changed

Lines changed: 215 additions & 11 deletions

File tree

packages/core/src/astgrep.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ export async function astGrepFindFiles(
1313
glob: ElementOrArray<string>,
1414
matcher: string | SgMatcher,
1515
options?: Omit<FindFilesOptions, "readText"> & CancellationOptions
16-
): ReturnType<AstGrep["search"]> {
16+
): ReturnType<Sg["search"]> {
1717
const { cancellationToken } = options || {}
1818
if (!glob) {
1919
throw new Error("glob is required")

packages/core/src/diff.ts

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -347,10 +347,12 @@ export function llmifyDiff(diff: string) {
347347
* @returns The diff as a string, with certain headers removed.
348348
*/
349349
export function createDiff(
350-
left: WorkspaceFile,
351-
right: WorkspaceFile,
350+
left: string | WorkspaceFile,
351+
right: string | WorkspaceFile,
352352
options?: { context?: number }
353353
) {
354+
if (typeof left === "string") left = { filename: "left", content: left }
355+
if (typeof right === "string") right = { filename: "right", content: right }
354356
const res = createTwoFilesPatch(
355357
left.filename || "",
356358
right.filename || "",

packages/core/src/promptcontext.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -332,7 +332,7 @@ export async function createPromptContext(
332332
await runtimeHost.python({ trace, ...(options || {}) }),
333333
teamsChannel: async (url) => createMicrosoftTeamsChannelClient(url),
334334
astGrep: async () =>
335-
Object.freeze<AstGrep>({
335+
Object.freeze<Sg>({
336336
search: (lang, glob, matcher) =>
337337
astGrepFindFiles(lang, glob, matcher, {
338338
cancellationToken,

packages/core/src/trace.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ import { parseTraceTree, TraceTree } from "./traceparser"
2828
import { fileCacheImage, fileWriteCached } from "./filecache"
2929
import { CancellationOptions } from "./cancellation"
3030
import { generateId } from "./id"
31+
import { createDiff } from "./diff"
3132

3233
export class TraceChunkEvent extends Event {
3334
constructor(
@@ -119,6 +120,15 @@ export class MarkdownTrace extends EventTarget implements OutputTrace {
119120
)
120121
}
121122

123+
diff(
124+
left: string | WorkspaceFile,
125+
right: string | WorkspaceFile,
126+
options?: { context?: number }
127+
) {
128+
const d = createDiff(left, right, options)
129+
this.fence(d, "diff")
130+
}
131+
122132
/**
123133
* Logs a markdown table
124134
* @param rows

packages/core/src/types/prompt_template.d.ts

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -859,6 +859,15 @@ interface OutputTrace extends ToolCallTrace {
859859
*/
860860
table(rows: object[]): void
861861

862+
/**
863+
* Computes and renders diff between two files.
864+
*/
865+
diff(
866+
left: string | WorkspaceFile,
867+
right: string | WorkspaceFile,
868+
options?: { context?: number }
869+
): void
870+
862871
/**
863872
* Logs a result item with a boolean value and a message.
864873
* @param value - The boolean value of the result item.
@@ -2334,8 +2343,8 @@ interface Parsers {
23342343
* Computes a diff between two files
23352344
*/
23362345
diff(
2337-
left: WorkspaceFile,
2338-
right: WorkspaceFile,
2346+
left: string | WorkspaceFile,
2347+
right: string | WorkspaceFile,
23392348
options?: DefDiffOptions
23402349
): string
23412350

@@ -4130,7 +4139,7 @@ interface SgRule {
41304139
/** A utility rule id and matches a node if the utility rule matches. */
41314140
matches?: string
41324141
}
4133-
interface SgRelation extends Rule {
4142+
interface SgRelation extends SgRule {
41344143
/**
41354144
* Specify how relational rule will stop relative to the target node.
41364145
*/
@@ -4189,7 +4198,7 @@ interface SgRoot {
41894198

41904199
type SgLang = OptionsOrString<"html" | "js" | "ts" | "tsx" | "css">
41914200

4192-
interface AstGrep {
4201+
interface Sg {
41934202
parse(file: WorkspaceFile, options: { lang?: SgLang }): Promise<SgRoot>
41944203
search(
41954204
lang: SgLang,
Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
script({
2+
title: "Generate TypeScript function documentation using AST insertion",
3+
accept: ".ts",
4+
files: "src/cowsay.ts",
5+
parameters: {
6+
applyEdits: {
7+
type: "boolean",
8+
default: false,
9+
description: "If true, the script will not modify the files.",
10+
},
11+
},
12+
})
13+
const { output } = env
14+
const { applyEdits } = env.vars
15+
const file = env.files[0]
16+
17+
// find all exported functions without comments
18+
const sg = await host.astGrep()
19+
const { matches, replace, commitEdits } = await sg.search("ts", file.filename, {
20+
rule: {
21+
kind: "export_statement",
22+
has: {
23+
kind: "function_declaration",
24+
not: {
25+
precedes: {
26+
kind: "comment",
27+
stopBy: "neighbor",
28+
},
29+
},
30+
},
31+
},
32+
})
33+
34+
// for each match, generate a docstring
35+
for (const match of matches) {
36+
const res = await runPrompt(
37+
(_) => {
38+
_.def("FILE", match.getRoot().root().text())
39+
_.def("FUNCTION", match.text())
40+
// this needs more eval-ing
41+
_.$`Generate a function documentation for <FUNCTION>.
42+
- Be concise. Use technical tone.
43+
- do NOT include types, this is for TypeScript.
44+
- Use docstring syntax.
45+
The full source of the file is in <FILE> for reference.`
46+
},
47+
{ model: "small", responseType: "text", label: match.text() }
48+
)
49+
// if generation is successful, insert the docs
50+
if (res.error) {
51+
output.warn(res.error.message)
52+
continue
53+
}
54+
const docs = docify(res.text.trim())
55+
const updated = `${docs}\n${match.text()}`
56+
replace(match, updated)
57+
}
58+
59+
// apply all edits and write to the file
60+
const modified = await commitEdits()
61+
if (applyEdits) {
62+
await workspace.writeFiles(modified)
63+
} else {
64+
output.diff(file, modified[0])
65+
output.warn(
66+
`edit not applied, use --vars 'applyEdits=true' to apply the edits`
67+
)
68+
}
69+
70+
// normalizes the docstring in case the LLM decides not to generate proper comments
71+
function docify(docs: string) {
72+
if (!/^\/\*\*.*.*\*\/$/s.test(docs))
73+
docs = `/**\n* ${docs.split(/\r?\n/g).join("\n* ")}\n*/`
74+
return docs.replace(/\n+$/, "")
75+
}

packages/sample/genaisrc/astgrep.genai.mjs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@ script({
44
model: "echo",
55
tests: {},
66
})
7-
const asg = await host.astGrep()
8-
const { matches } = await asg.search("ts", "src/*.ts", "console.log($META)")
7+
const sg = await host.astGrep()
8+
const { matches } = await sg.search("ts", "src/*.ts", "console.log($META)")
99
if (matches.length < 2) throw new Error("No matches src/*.ts found")
1010
for (const match of matches) {
1111
const t = match.text()
@@ -17,7 +17,7 @@ const {
1717
matches: matches2,
1818
replace,
1919
commitEdits,
20-
} = await asg.search("ts", "src/fib.ts", {
20+
} = await sg.search("ts", "src/fib.ts", {
2121
rule: {
2222
kind: "function_declaration",
2323
not: {

packages/sample/src/cowsay.ts

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
interface CowsayOptions {
2+
text: string
3+
mode?: "say" | "think"
4+
eyes?: string
5+
tongue?: string
6+
}
7+
8+
export function cowsay(options: CowsayOptions | string): string {
9+
// Handle string argument
10+
const opts: CowsayOptions =
11+
typeof options === "string" ? { text: options } : options
12+
13+
// Default options
14+
const { text = "", mode = "say", eyes = "oo", tongue = " " } = opts
15+
16+
// Split text into lines
17+
const lines = formatText(text)
18+
19+
// Create the speech bubble
20+
const bubble = createBubble(lines, mode)
21+
22+
// Create the cow
23+
const cow = createCow(eyes, tongue, mode)
24+
25+
// Combine the bubble and cow
26+
return bubble + cow
27+
}
28+
29+
function formatText(text: string, maxWidth: number = 40): string[] {
30+
if (!text) return [""]
31+
32+
const words = text.split(" ")
33+
const lines: string[] = []
34+
let currentLine = ""
35+
36+
for (const word of words) {
37+
if (currentLine.length + word.length + 1 <= maxWidth) {
38+
currentLine += (currentLine ? " " : "") + word
39+
} else {
40+
lines.push(currentLine)
41+
currentLine = word
42+
}
43+
}
44+
45+
if (currentLine) {
46+
lines.push(currentLine)
47+
}
48+
49+
return lines
50+
}
51+
52+
function createBubble(lines: string[], mode: "say" | "think"): string {
53+
if (lines.length === 0) return ""
54+
55+
const maxLength = Math.max(...lines.map((line) => line.length))
56+
let result = " " + "_".repeat(maxLength + 2) + "\n"
57+
58+
if (lines.length === 1) {
59+
const line = lines[0]
60+
const padding = " ".repeat(maxLength - line.length)
61+
result +=
62+
mode === "say"
63+
? `< ${line}${padding} >\n`
64+
: `( ${line}${padding} )\n`
65+
} else {
66+
lines.forEach((line, i) => {
67+
const padding = " ".repeat(maxLength - line.length)
68+
let prefix, suffix
69+
70+
if (i === 0) {
71+
prefix = mode === "say" ? "/ " : "( "
72+
suffix = mode === "say" ? " \\" : " )"
73+
} else if (i === lines.length - 1) {
74+
prefix = mode === "say" ? "\\ " : "( "
75+
suffix = mode === "say" ? " /" : " )"
76+
} else {
77+
prefix = mode === "say" ? "| " : "( "
78+
suffix = mode === "say" ? " |" : " )"
79+
}
80+
81+
result += `${prefix}${line}${padding}${suffix}\n`
82+
})
83+
}
84+
85+
result += " " + "-".repeat(maxLength + 2) + "\n"
86+
return result
87+
}
88+
89+
/**
90+
* Create the ASCII cow
91+
*/
92+
function createCow(
93+
eyes: string,
94+
tongue: string,
95+
mode: "say" | "think"
96+
): string {
97+
return ` \\ ^__^
98+
\\ (${eyes})\\_______
99+
(__)\\ )\\/\\
100+
${tongue}||----w |
101+
|| ||
102+
`
103+
}
104+
105+
export function cowthink(options: CowsayOptions | string): string {
106+
const opts = typeof options === "string" ? { text: options } : options
107+
return cowsay({ ...opts, mode: "think" })
108+
}

0 commit comments

Comments
 (0)