v0.0.3v0.0.3

1 files changed, 38 insertions, 22 deletions

diff --git a/src/extension.ts b/src/extension.ts
index 3a1de6b..3bc0995 100644
--- a/src/extension.ts
+++ b/src/extension.ts
@@ -1,9 +1,7 @@
-// The module 'vscode' contains the VS Code extensibility API
-// Import the module and reference it with the alias vscode in your code below
 import * as vscode from 'vscode';
 
-// this method is called when your extension is activated
-// your extension is activated the very first time the command is executed
+const fnRegex = /^(\t*)(.*)err(\s?):=(.+?)$/
+
 export function activate(context: vscode.ExtensionContext) {
 	context.subscriptions.push(vscode.languages.registerCodeActionsProvider('go', new ErrorsWrapper(), {
 		providedCodeActionKinds: ErrorsWrapper.providedCodeActionKinds
@@ -11,28 +9,37 @@ export function activate(context: vscode.ExtensionContext) {
 	context.subscriptions.push(
 		vscode.commands.registerCommand('gotools.wrap-error', () => {
 			const editor = vscode.window.activeTextEditor;
-
 			if (editor) {
 				const document = editor.document;
-				const selection = editor.selection;
-				const line = document.lineAt(editor.selection.start.line)
-				const idx = line.firstNonWhitespaceCharacterIndex
-				const word = document.getText(selection);
-				const space = "\t".repeat(idx)
-				let errPrefix = ""
-				if (word.indexOf("=") == -1) {
-					errPrefix = "err :="
+				const line = document.lineAt(editor.selection.start.line);
+				const matches = line.text.match(fnRegex);
+				if (matches == null) {
+					return;
+				}
+				if (matches.length > 0) {
+					const intendation = matches[1];
+					const extravars = matches[2].split(',').map(x => x.trim()).filter(x => x);
+					const rest = matches[4].trim();
+					editor.edit(editBuilder => {
+						if (extravars.filter(x => x != "_").length > 0) {
+							editBuilder.insert(
+								new vscode.Position(line.lineNumber + 1, 0),
+								`${intendation}if err != nil {\n${intendation}\treturn err\n${intendation}}`
+							);
+						} else {
+							extravars.push("err");
+							editBuilder.replace(
+								line.range,
+								`${intendation}if ${extravars.join(", ")} := ${rest}; err != nil {\n${intendation}\treturn err\n${intendation}}`
+							);
+						}
+					});
 				}
-				editor.edit(editBuilder => {
-					editBuilder.replace(selection, `if ${errPrefix}${word}; err != nil {\n${space}\rreturn err\n${space}}`);
-				});
 			}
 		})
 	);
-
 }
 
-// this method is called when your extension is deactivated
 export function deactivate() { }
 
 export class ErrorsWrapper implements vscode.CodeActionProvider {
@@ -42,10 +49,19 @@ export class ErrorsWrapper implements vscode.CodeActionProvider {
 	];
 
 	public provideCodeActions(document: vscode.TextDocument, range: vscode.Range): vscode.CodeAction[] | undefined {
-		const action = new vscode.CodeAction('Wrap error', vscode.CodeActionKind.RefactorRewrite);
-		action.command = { command: 'gotools.wrap-error', title: 'Wrap with error block', tooltip: '' };
+		const editor = vscode.window.activeTextEditor;
+		if (!editor) {
+			return undefined;
+		}
+		const selection = editor.selection;
+		const line = document.lineAt(editor.selection.start.line);
+		if (!fnRegex.test(line.text)) {
+			return undefined;
+		}
+		const action = new vscode.CodeAction('Add error checking', vscode.CodeActionKind.RefactorRewrite);
+		action.command = { command: 'gotools.wrap-error', title: 'Add error checking block', tooltip: '' };
 		return [
-			action
-		]
+			action,
+		];
 	}
 }