src/goGenerateMethod: make method generation more automatic using code action

sslime336 · sslime336 · commit 3623adc82b90 · 2023-02-16T17:12:38.000+08:00
diff --git a/package.json b/package.json
@@ -353,6 +353,11 @@
         "title": "Go: Generate Unit Tests For Function",
         "description": "Generates unit tests for the selected function in the current file"
       },
+      {
+        "command": "go.generate.method",
+        "title": "Go: Generate Method",
+        "description": "Generates method for the given type"
+      },
       {
         "command": "go.impl.cursor",
         "title": "Go: Generate Interface Stubs",
diff --git a/src/goGenerateMethod.ts b/src/goGenerateMethod.ts
@@ -0,0 +1,119 @@
+import vscode = require('vscode');
+import { CommandFactory } from './commands';
+
+const CommandTitle = 'Generate method';
+const Command = 'go.generate.method';
+
+export const goGenerateMethod: CommandFactory = () => (
+	uncommontypeName: string,
+	needPtrReceiver: boolean,
+	endPos: vscode.Position
+) => {
+	const editor = vscode.window.activeTextEditor;
+	if (!editor) {
+		vscode.window.showErrorMessage('No active editor found.');
+		return;
+	}
+	const receiverName = getReceiverName(uncommontypeName);
+	let methodTpl = `\n\nfunc ($\{1:${receiverName}} *${uncommontypeName}) $\{2:methodName}($3) $4 {\n\t$5\n}$0`;
+	if (!needPtrReceiver) {
+		methodTpl = `\n\nfunc ($\{1:${receiverName}} ${uncommontypeName}) $\{2:methodName}($3) $4 {\n\t$5\n}$0`;
+	}
+	editor.insertSnippet(new vscode.SnippetString(methodTpl), endPos);
+};
+
+export class MethodGenerationProvider implements vscode.CodeActionProvider {
+	public static readonly providedCodeActionKinds = [vscode.CodeActionKind.Refactor];
+
+	async provideCodeActions(
+		document: vscode.TextDocument,
+		range: vscode.Range | vscode.Selection,
+		_context: vscode.CodeActionContext,
+		_token: vscode.CancellationToken
+	): Promise<vscode.CodeAction[] | undefined> {
+		const lineText = document.lineAt(range.start.line).text;
+		// TODO(sslime336): support the uncommontypes defined in type block.
+		const uncommontypeName = await this.getUncommontypeName(lineText);
+
+		if (uncommontypeName === '') {
+			return;
+		}
+
+		let documentSymbols: vscode.DocumentSymbol[] = [];
+		await vscode.commands
+			.executeCommand<vscode.DocumentSymbol[]>('vscode.executeDocumentSymbolProvider', document.uri)
+			.then((symbols) => {
+				documentSymbols = symbols.filter((symbol) => {
+					const res = symbol.name === uncommontypeName;
+					return res;
+				});
+			});
+
+		if (documentSymbols.length === 0) {
+			return;
+		}
+
+		const endPos = documentSymbols[0].range.end;
+
+		const genPointerReceiverMethod = new vscode.CodeAction(
+			'generate method with pointer receiver',
+			vscode.CodeActionKind.Refactor
+		);
+		genPointerReceiverMethod.command = {
+			title: CommandTitle,
+			command: Command,
+			arguments: [uncommontypeName, true, endPos]
+		};
+
+		const genValueReceiverMethod = new vscode.CodeAction(
+			'generate method with value receiver',
+			vscode.CodeActionKind.Refactor
+		);
+		genValueReceiverMethod.command = {
+			title: CommandTitle,
+			command: Command,
+			arguments: [uncommontypeName, false, endPos]
+		};
+
+		return [genPointerReceiverMethod, genValueReceiverMethod];
+	}
+
+	resolveCodeAction?(
+		_codeAction: vscode.CodeAction,
+		_token: vscode.CancellationToken
+	): vscode.ProviderResult<vscode.CodeAction> {
+		return;
+	}
+
+	// getUncommontypeName returns the user defined type's name from type definitions
+	// that starts with single keyword type.
+	// The types defined in the type definition block will not satisfied the regexp.
+	private async getUncommontypeName(lineText: string): Promise<string> {
+		const regexp = /type ([^0-9]\w+) [^0-9]\w+/;
+		const matches = lineText.match(regexp);
+		if (!matches) {
+			return '';
+		}
+		return matches[1];
+	}
+}
+
+// getReceiverName returns the default receiver name which is constructed with uppercase
+// characters picked from the structName, it will return the first character in lowercase if
+// the structName contains no uppercased character.
+function getReceiverName(structName: string): string {
+	let res = '';
+	structName
+		.split('')
+		.filter((ch) => isUpperCase(ch))
+		.forEach((ch) => (res += ch.toLowerCase()));
+	if (res === '') {
+		res = structName.charAt(0);
+	}
+	return res;
+}
+
+const isUpperCase = (ch: string): boolean => {
+	const c = ch.charCodeAt(0);
+	return 65 <= c && c <= 90;
+};
diff --git a/src/goMain.ts b/src/goMain.ts
@@ -66,6 +66,7 @@ import { GoExplorerProvider } from './goExplorer';
 import { GoExtensionContext } from './context';
 import * as commands from './commands';
 import { toggleVulncheckCommandFactory, VulncheckOutputLinkProvider } from './goVulncheck';
+import { MethodGenerationProvider, goGenerateMethod } from './goGenerateMethod';
 
 const goCtx: GoExtensionContext = {};
 
@@ -132,6 +133,12 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<ExtensionA
 	goCtx.vetDiagnosticCollection = vscode.languages.createDiagnosticCollection('go-vet');
 	ctx.subscriptions.push(goCtx.vetDiagnosticCollection);
 
+	const genStructMethodProvidor = vscode.languages.registerCodeActionsProvider('go', new MethodGenerationProvider(), {
+		providedCodeActionKinds: MethodGenerationProvider.providedCodeActionKinds
+	});
+	ctx.subscriptions.push(genStructMethodProvidor);
+
+	registerCommand('go.generate.method', goGenerateMethod);
 	registerCommand('go.gopath', commands.getCurrentGoPath);
 	registerCommand('go.goroot', commands.getCurrentGoRoot);
 	registerCommand('go.locate.tools', commands.getConfiguredGoTools);
