diff --git a/src/runtime/generator.ts b/src/runtime/generator.ts index 0f7598e..b2f43dd 100644 --- a/src/runtime/generator.ts +++ b/src/runtime/generator.ts @@ -16,6 +16,18 @@ const getGeneratorContext = (cwd: string): Fig.GeneratorContext => { }; }; +// Matches Fig's cleanOutput: CRLF becomes LF, the cursor-show sequence and leading/trailing newlines are removed. +const cleanOutput = (output: string): string => output.replaceAll("\r\n", "\n").replaceAll("\u001b[?25h", "").replace(/^\n+/, "").replace(/\n+$/, ""); + +// Custom generators call executeShellCommand themselves. Like Fig, run their commands in the generator's +// working directory unless they pass one, and give them cleaned output (e.g. `npm prefix` without a trailing newline). +const buildCustomExecuteCommand = + (executeShellCommand: Fig.ExecuteCommandFunction, cwd: string): Fig.ExecuteCommandFunction => + async (input) => { + const output = await executeShellCommand({ ...input, cwd: input.cwd ?? cwd }); + return { ...output, stdout: cleanOutput(output.stdout), stderr: cleanOutput(output.stderr) }; + }; + export const executeGenerator = async (generator: Fig.Generator, tokens: string[], cwd: string, signal?: AbortSignal): Promise => { const generatorTiming = startTiming(); try { @@ -39,7 +51,7 @@ export const executeGenerator = async (generator: Fig.Generator, tokens: string[ } if (custom) { - suggestions.push(...(await custom(tokens, executeShellCommand, getGeneratorContext(cwd)))); + suggestions.push(...(await custom(tokens, buildCustomExecuteCommand(executeShellCommand, cwd), getGeneratorContext(cwd)))); } if (template != null) { diff --git a/src/tests/runtime/generator.test.ts b/src/tests/runtime/generator.test.ts index e240fbc..27f602c 100644 --- a/src/tests/runtime/generator.test.ts +++ b/src/tests/runtime/generator.test.ts @@ -4,7 +4,10 @@ import { jest } from "@jest/globals"; import { clearGeneratorState, clearImplicitGeneratorState, getGeneratorSuggestions } from "../../runtime/generatorCache.js"; -import { getGeneratorQueryTerm } from "../../runtime/generator.js"; +import { executeGenerator, getGeneratorQueryTerm } from "../../runtime/generator.js"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; const deferred = () => { let resolve!: (value: T) => void; @@ -210,3 +213,38 @@ test("reuses a generator when a path query moves past a delimiter", async () => expect(calls).toBe(1); }); + +describe("custom generator executeShellCommand", () => { + const nodeGenerator = (script: string, cwd?: string): Fig.Generator => ({ + custom: async (_tokens, executeShellCommand) => { + const { stdout } = await executeShellCommand({ command: process.execPath, args: ["-e", script], cwd }); + return [{ name: stdout }]; + }, + }); + + test("runs in the generator's cwd when the generator does not pass one", async () => { + const dir = fs.realpathSync(fs.mkdtempSync(path.join(os.tmpdir(), "generator-cwd-"))); + try { + const [suggestion] = await executeGenerator(nodeGenerator("process.stdout.write(process.cwd())"), [], dir); + expect(suggestion.name).toBe(dir); + } finally { + fs.rmSync(dir, { recursive: true, force: true }); + } + }); + + test("keeps a cwd passed by the generator", async () => { + const dir = fs.realpathSync(fs.mkdtempSync(path.join(os.tmpdir(), "generator-cwd-"))); + try { + const [suggestion] = await executeGenerator(nodeGenerator("process.stdout.write(process.cwd())", dir), [], os.tmpdir()); + expect(suggestion.name).toBe(dir); + } finally { + fs.rmSync(dir, { recursive: true, force: true }); + } + }); + + test("cleans output like Fig: CRLF to LF, no leading or trailing newlines", async () => { + const script = String.raw`process.stdout.write("\r\n\n a\r\nb \r\n\n")`; + const [suggestion] = await executeGenerator(nodeGenerator(script), [], process.cwd()); + expect(suggestion.name).toBe(" a\nb "); + }); +});