From c049f23a0c5608fce11aa74b048b0a65c3d027dc Mon Sep 17 00:00:00 2001 From: Zhiming Ma Date: Tue, 28 Nov 2023 16:57:02 +0800 Subject: [PATCH] feat(agent): format indentation if not match with editor config. (#911) --- clients/tabby-agent/src/CompletionContext.ts | 3 + .../src/postprocess/formatIndentation.test.ts | 166 ++++++++++++++++++ .../src/postprocess/formatIndentation.ts | 100 +++++++++++ clients/tabby-agent/src/postprocess/index.ts | 2 + clients/vscode/src/TabbyCompletionProvider.ts | 16 ++ 5 files changed, 287 insertions(+) create mode 100644 clients/tabby-agent/src/postprocess/formatIndentation.test.ts create mode 100644 clients/tabby-agent/src/postprocess/formatIndentation.ts diff --git a/clients/tabby-agent/src/CompletionContext.ts b/clients/tabby-agent/src/CompletionContext.ts index 954931b..0bcb5ed 100644 --- a/clients/tabby-agent/src/CompletionContext.ts +++ b/clients/tabby-agent/src/CompletionContext.ts @@ -6,6 +6,7 @@ export type CompletionRequest = { language: string; text: string; position: number; + indentation?: string; clipboard?: string; manually?: boolean; }; @@ -36,6 +37,7 @@ function isAtLineEndExcludingAutoClosedChar(suffix: string) { export class CompletionContext { filepath: string; language: string; + indentation?: string; text: string; position: number; @@ -57,6 +59,7 @@ export class CompletionContext { this.language = request.language; this.text = request.text; this.position = request.position; + this.indentation = request.indentation; this.prefix = request.text.slice(0, request.position); this.suffix = request.text.slice(request.position); diff --git a/clients/tabby-agent/src/postprocess/formatIndentation.test.ts b/clients/tabby-agent/src/postprocess/formatIndentation.test.ts new file mode 100644 index 0000000..9a35447 --- /dev/null +++ b/clients/tabby-agent/src/postprocess/formatIndentation.test.ts @@ -0,0 +1,166 @@ +import { expect } from "chai"; +import { documentContext, inline } from "./testUtils"; +import { formatIndentation } from "./formatIndentation"; + +describe("postprocess", () => { + describe("formatIndentation", () => { + it("should format indentation if first line of completion is over indented.", () => { + const context = { + ...documentContext` + function clamp(n: number, max: number, min: number): number { + ║ + } + `, + indentation: " ", + language: "typescript", + }; + const completion = inline` + ├ return Math.max(Math.min(n, max), min);┤ + `; + const expected = inline` + ├return Math.max(Math.min(n, max), min);┤ + `; + expect(formatIndentation(context)(completion)).to.eq(expected); + }); + + it("should format indentation if first line of completion is wrongly indented.", () => { + const context = { + ...documentContext` + function clamp(n: number, max: number, min: number): number { + ║ + } + `, + indentation: " ", + language: "typescript", + }; + const completion = inline` + ├ return Math.max(Math.min(n, max), min);┤ + `; + const expected = inline` + ├ return Math.max(Math.min(n, max), min);┤ + `; + expect(formatIndentation(context)(completion)).to.eq(expected); + }); + + it("should format indentation if completion lines is over indented.", () => { + const context = { + ...documentContext` + def findMax(arr):║ + `, + indentation: " ", + language: "python", + }; + const completion = inline` + ├ + max = arr[0] + for i in range(1, len(arr)): + if arr[i] > max: + max = arr[i] + return max + }┤ + `; + const expected = inline` + ├ + max = arr[0] + for i in range(1, len(arr)): + if arr[i] > max: + max = arr[i] + return max + }┤ + `; + expect(formatIndentation(context)(completion)).to.eq(expected); + }); + + it("should format indentation if completion lines is wrongly indented.", () => { + const context = { + ...documentContext` + def findMax(arr):║ + `, + indentation: " ", + language: "python", + }; + const completion = inline` + ├ + max = arr[0] + for i in range(1, len(arr)): + if arr[i] > max: + max = arr[i] + return max + }┤ + `; + const expected = inline` + ├ + max = arr[0] + for i in range(1, len(arr)): + if arr[i] > max: + max = arr[i] + return max + }┤ + `; + expect(formatIndentation(context)(completion)).to.eq(expected); + }); + + it("should keep it unchanged if it no indentation specified.", () => { + const context = { + ...documentContext` + def findMax(arr):║ + `, + indentation: undefined, + language: "python", + }; + const completion = inline` + ├ + max = arr[0] + for i in range(1, len(arr)): + if arr[i] > max: + max = arr[i] + return max + }┤ + `; + expect(formatIndentation(context)(completion)).to.eq(completion); + }); + + it("should keep it unchanged if there is indentation in the context.", () => { + const context = { + ...documentContext` + def hello(): + return "world" + + def findMax(arr):║ + `, + indentation: "\t", + language: "python", + }; + const completion = inline` + ├ + max = arr[0] + for i in range(1, len(arr)): + if arr[i] > max: + max = arr[i] + return max + }┤ + `; + expect(formatIndentation(context)(completion)).to.eq(completion); + }); + + it("should keep it unchanged if it is well indented.", () => { + const context = { + ...documentContext` + def findMax(arr):║ + `, + indentation: " ", + language: "python", + }; + const completion = inline` + ├ + max = arr[0] + for i in range(1, len(arr)): + if arr[i] > max: + max = arr[i] + return max + }┤ + `; + expect(formatIndentation(context)(completion)).to.eq(completion); + }); + }); +}); diff --git a/clients/tabby-agent/src/postprocess/formatIndentation.ts b/clients/tabby-agent/src/postprocess/formatIndentation.ts new file mode 100644 index 0000000..0761038 --- /dev/null +++ b/clients/tabby-agent/src/postprocess/formatIndentation.ts @@ -0,0 +1,100 @@ +import { CompletionContext } from "../Agent"; +import { PostprocessFilter, logger } from "./base"; +import { isBlank, splitLines } from "../utils"; + +function detectIndentation(lines: string[]): string | null { + const matches = { + "\t": 0, + " ": 0, + " ": 0, + }; + for (const line of lines) { + if (line.match(/^\t/)) { + matches["\t"]++; + } else { + const spaces = line.match(/^ */)[0].length; + if (spaces > 0) { + if (spaces % 4 === 0) { + matches[" "]++; + } + if (spaces % 2 === 0) { + matches[" "]++; + } + } + } + } + if (matches["\t"] > 0) { + return "\t"; + } + if (matches[" "] > matches[" "]) { + return " "; + } + if (matches[" "] > 0) { + return " "; + } + return null; +} + +function getIndentLevel(line: string, indentation: string): number { + if (indentation === "\t") { + return line.match(/^\t*/g)[0].length; + } else { + const spaces = line.match(/^ */)[0].length; + return spaces / indentation.length; + } +} + +export function formatIndentation(context: CompletionContext): PostprocessFilter { + return (input) => { + const { prefixLines, suffixLines, indentation } = context; + const inputLines = splitLines(input); + + // if no indentation is specified + if (!indentation) { + return input; + } + + // if there is any indentation in context, the server output should have learned from it + const prefixLinesForDetection = isBlank(prefixLines[prefixLines.length - 1]) + ? prefixLines.slice(0, prefixLines.length - 1) + : prefixLines; + if (prefixLines.length > 1 && detectIndentation(prefixLinesForDetection) !== null) { + return input; + } + const suffixLinesForDetection = suffixLines.slice(1); + if (suffixLines.length > 1 && detectIndentation(suffixLinesForDetection) !== null) { + return input; + } + + // if the input is well indented with specific indentation + const inputLinesForDetection = inputLines.map((line, index) => { + return index === 0 ? prefixLines[prefixLines.length - 1] + line : line; + }); + const inputIndentation = detectIndentation(inputLinesForDetection); + if (inputIndentation === null || inputIndentation === indentation) { + return input; + } + + // otherwise, do formatting + const formatted = inputLinesForDetection.map((line, index) => { + const level = getIndentLevel(inputLinesForDetection[index], inputIndentation); + if (level === 0) { + return inputLines[index]; + } + const rest = line.slice(inputIndentation.length * level); + if (index === 0) { + // for first line + if (!isBlank(prefixLines[prefixLines.length - 1])) { + return inputLines[0]; + } else { + return indentation.repeat(level).slice(prefixLines[prefixLines.length - 1].length) + rest; + } + } else { + // for next lines + return indentation.repeat(level) + rest; + } + }); + logger.debug({ prefixLines, suffixLines, inputLines, formatted }, "Format indentation."); + return formatted.join(""); + }; +} diff --git a/clients/tabby-agent/src/postprocess/index.ts b/clients/tabby-agent/src/postprocess/index.ts index 0814c81..ad5a62c 100644 --- a/clients/tabby-agent/src/postprocess/index.ts +++ b/clients/tabby-agent/src/postprocess/index.ts @@ -6,6 +6,7 @@ import { removeRepetitiveBlocks } from "./removeRepetitiveBlocks"; import { removeRepetitiveLines } from "./removeRepetitiveLines"; import { removeLineEndsWithRepetition } from "./removeLineEndsWithRepetition"; import { limitScope } from "./limitScope"; +import { formatIndentation } from "./formatIndentation"; import { trimSpace } from "./trimSpace"; import { dropDuplicated } from "./dropDuplicated"; import { dropBlank } from "./dropBlank"; @@ -33,6 +34,7 @@ export async function postCacheProcess( .then(applyFilter(removeRepetitiveBlocks(context), context)) .then(applyFilter(removeRepetitiveLines(context), context)) .then(applyFilter(limitScope(context, config["limitScope"]), context)) + .then(applyFilter(formatIndentation(context), context)) .then(applyFilter(dropDuplicated(context), context)) .then(applyFilter(trimSpace(context), context)) .then(applyFilter(dropBlank(), context)); diff --git a/clients/vscode/src/TabbyCompletionProvider.ts b/clients/vscode/src/TabbyCompletionProvider.ts index e08be1f..a317c67 100644 --- a/clients/vscode/src/TabbyCompletionProvider.ts +++ b/clients/vscode/src/TabbyCompletionProvider.ts @@ -72,6 +72,7 @@ export class TabbyCompletionProvider extends EventEmitter implements InlineCompl language: document.languageId, // https://code.visualstudio.com/docs/languages/identifiers text: additionalContext.prefix + document.getText() + additionalContext.suffix, position: additionalContext.prefix.length + document.offsetAt(position), + indentation: this.getEditorIndentation(), clipboard: await env.clipboard.readText(), manually: context.triggerKind === InlineCompletionTriggerKind.Invoke, }; @@ -166,6 +167,21 @@ export class TabbyCompletionProvider extends EventEmitter implements InlineCompl } } + private getEditorIndentation(): string | undefined { + const editor = window.activeTextEditor; + if (!editor) { + return undefined; + } + + const { insertSpaces, tabSize } = editor.options; + if (insertSpaces && typeof tabSize === "number" && tabSize > 0) { + return " ".repeat(tabSize); + } else if (!insertSpaces) { + return "\t"; + } + return undefined; + } + private updateConfiguration() { if (!workspace.getConfiguration("editor").get("inlineSuggest.enabled", true)) { this.triggerMode = "disabled";