Introduce Kedro viz in VSCode Extension

.gitignore

@@ -9,4 +9,4 @@ node_modules
 bundled/libs/
 **/__pycache__
 **/.pytest_cache
-**/.vs
+**/.vs

.vscodeignore

@@ -1,7 +1,7 @@
 .vscode/**
 .vscode-test/**
 out/**
-node_modules/**
+**/node_modules/**
 src/**
 .gitignore
 .yarnrc

bundled/tool/lsp_server.py

@@ -549,7 +549,8 @@ def definition_from_flowchart(ls, word):
     It will `block` the main thread, which can be tested by trying to show
     completion items.
     """
-    result = definition(LSP_SERVER, params=None, word="companies")
+    word = word[0]
+    result = definition(LSP_SERVER, params=None, word=word)
     return result
 
 ### End of  kedro-lsp

package.json

@@ -2,7 +2,7 @@
     "name": "Kedro",
     "displayName": "Kedro",
     "description": "A Kedro VSCode Extension.",
-    "version": "0.1.0",
+    "version": "0.2.0-rc0",
     "preview": false,
     "serverInfo": {
         "name": "Kedro",
@@ -55,26 +55,28 @@
     "scripts": {
         "vscode:prepublish": "npm run package",
         "compile": "webpack",
-        "watch": "webpack --watch",
+        "watch": "concurrently \"npm:watch --prefix ./webview\" \"webpack --watch\"",
         "package": "webpack --mode production --devtool source-map --config ./webpack.config.js",
         "compile-tests": "tsc -p . --outDir out",
         "watch-tests": "tsc -p . -w --outDir out",
         "pretest": "npm run compile-tests && npm run compile && npm run lint",
         "lint": "eslint src --ext ts",
         "format-check": "prettier --check 'src/**/*.ts' 'build/**/*.yml' '.github/**/*.yml'",
         "test": "node ./out/test/runTest.js",
-        "vsce-package": "vsce package -o kedro.vsix"
+        "vsce-package": "vsce package -o kedro.vsix",
+        "build:webview": "npm run build --prefix ./webview",
+        "install:webview": "npm install --prefix ./webview"
     },
     "contributes": {
         "icons": {
-        "kedro-logo": {
-            "description": "Kedro icon",
-            "default": {
-            "fontPath": "assets/kedro-logo/fonts/kedro-logo.woff",
-            "fontCharacter": "\\e900"
+            "kedro-logo": {
+                "description": "Kedro icon",
+                "default": {
+                    "fontPath": "assets/kedro-logo/fonts/kedro-logo.woff",
+                    "fontCharacter": "\\e900"
+                }
             }
-        }
-    },
+        },
         "configuration": {
             "properties": {
                 "yaml.schemas": {
@@ -160,29 +162,36 @@
                 "title": "Send Definition Request",
                 "category": "kedro",
                 "command": "kedro.sendDefinitionRequest"
+            },
+            {
+                "command": "kedro.runKedroViz",
+                "category": "kedro",
+                "title": "Run Kedro Viz"
             }
 
         ]
     },
     "dependencies": {
         "@vscode/python-extension": "^1.0.5",
         "fs-extra": "^11.2.0",
+        "node-fetch": "^3.3.2",
         "vscode-languageclient": "^8.1.0"
     },
     "devDependencies": {
         "@types/fs-extra": "^11.0.4",
-        "@types/vscode": "1.78.0",
         "@types/glob": "^8.1.0",
         "@types/node": "16.x",
+        "@types/vscode": "1.78.0",
         "@typescript-eslint/eslint-plugin": "^7.14.1",
         "@typescript-eslint/parser": "^7.11.0",
         "@vscode/test-electron": "^2.3.8",
         "@vscode/vsce": "^2.22.0",
+        "concurrently": "^8.2.2",
         "eslint": "^8.56.0",
         "glob": "^10.3.10",
         "prettier": "^3.1.1",
-        "typescript": "^5.3.3",
         "ts-loader": "^9.5.1",
+        "typescript": "^5.3.3",
         "webpack": "^5.89.0",
         "webpack-cli": "^5.1.4"
     }

src/common/commands.ts

@@ -60,7 +60,7 @@ export async function executeServerCommand(lsClient: LanguageClient | undefined)
     logger.info(`${commandName} result: ${JSON.stringify(result, undefined, 2)}`);
 }
 
-export async function executeServerDefinitionCommand(lsClient: LanguageClient | undefined) {
+export async function executeServerDefinitionCommand(lsClient: LanguageClient | undefined, word?: string | undefined) {
     if (!lsClient || lsClient.state !== State.Running) {
         await vscode.window.showErrorMessage('There is no language server running.');
         return;
@@ -71,28 +71,28 @@ export async function executeServerDefinitionCommand(lsClient: LanguageClient |
     }
 
     const commandName = 'kedro.goToDefinitionFromFlowchart';
+    let target: any = word;
+    if (!target) {
+        target = await window.showInputBox({
+            placeHolder: 'Type the name of the dataset/parameters, i.e. companies',
+        });
+    }
 
     logger.info(`executing command: '${commandName}'`);
 
     const result: any[] | undefined = await vscode.commands.executeCommand(
         commandName /* if your command accepts arguments you can pass them here */,
+        target,
     );
     logger.info(`${commandName} result: ${JSON.stringify(result, undefined, 2)}`);
     if (result && result.length > 0) {
         const location = result[0];
         const uri: vscode.Uri = vscode.Uri.parse(location.uri);
         const range = location.range;
 
-        vscode.window.showTextDocument(uri,
-            {
-                selection: range,
-                viewColumn: vscode.ViewColumn.One,
-            }
-        );
-
+        vscode.window.showTextDocument(uri, {
+            selection: range,
+            viewColumn: vscode.ViewColumn.One,
+        });
+    }
 }
-
-
-
-
-

src/common/utilities.ts

@@ -6,6 +6,8 @@ import * as path from 'path';
 import { LogLevel, Uri, WorkspaceFolder } from 'vscode';
 import { Trace } from 'vscode-jsonrpc/node';
 import { getWorkspaceFolders } from './vscodeapi';
+import fetch from 'node-fetch';
+import KedroVizPanel from '../webview/vizWebView';
 
 function logLevelToTrace(logLevel: LogLevel): Trace {
     switch (logLevel) {
@@ -65,3 +67,10 @@ export async function getProjectRoot(): Promise<WorkspaceFolder> {
         return rootWorkspace;
     }
 }
+
+export async function fetchAndUpdateProjectData(): Promise<void> {
+    fetch('http://127.0.0.1:3131/api/main')
+        .then((response: { text: () => any }) => response.text())
+        .then((data: string) => KedroVizPanel.currentPanel?.updateData(data))
+        .catch((err: { message: string }) => console.error('Error: ' + err.message));
+}

src/extension.ts

@@ -1,6 +1,6 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
-import { selectEnvironment, executeServerCommand,  executeServerDefinitionCommand } from './common/commands';
+import { selectEnvironment, executeServerCommand, executeServerDefinitionCommand } from './common/commands';
 import * as vscode from 'vscode';
 import { LanguageClient } from 'vscode-languageclient/node';
 import { registerLogger, traceError, traceLog, traceVerbose } from './common/log/logging';
@@ -22,8 +22,12 @@ import {
 import { loadServerDefaults } from './common/setup';
 import { getLSClientTraceLevel, getProjectRoot } from './common/utilities';
 import { createOutputChannel, onDidChangeConfiguration, registerCommand } from './common/vscodeapi';
+import KedroVizPanel from './webview/vizWebView';
+import { runKedroVizServer } from './webview/vizServer';
 
 let lsClient: LanguageClient | undefined;
+let kedroVizProcess: any;
+
 export async function activate(context: vscode.ExtensionContext): Promise<void> {
     // This is required to get server name and module. This should be
     // the first thing that we do in this extension.
@@ -53,6 +57,19 @@ export async function activate(context: vscode.ExtensionContext): Promise<void>
         }),
     );
 
+    context.subscriptions.push(
+        vscode.commands.registerCommand('kedro.runKedroViz', () => {
+            KedroVizPanel.createOrShow(context.extensionUri, lsClient);
+        }),
+    );
+    context.subscriptions.push(
+        vscode.commands.registerCommand('kedro.updateTheme', () => {
+            if (KedroVizPanel.currentPanel) {
+                KedroVizPanel.currentPanel.updateTheme();
+            }
+        }),
+    );
+
     // Log Server information
     traceLog(`Name: ${serverInfo.name}`);
     traceLog(`Module: ${serverInfo.module}`);
@@ -77,6 +94,10 @@ export async function activate(context: vscode.ExtensionContext): Promise<void>
         console.log('===============DEBUG============');
         console.log(interpreterDetails);
         console.log('===============DEBUG============');
+
+        // Start kedro viz server
+        kedroVizProcess = await runKedroVizServer();
+
         if (interpreterDetails.path) {
             traceVerbose(`Using interpreter from Python extension: ${interpreterDetails.path.join(' ')}`);
             lsClient = await restartServer(serverId, serverName, outputChannel, lsClient, env);
@@ -85,9 +106,9 @@ export async function activate(context: vscode.ExtensionContext): Promise<void>
 
         traceError(
             'Python interpreter missing:\r\n' +
-            '[Option 1] Select python interpreter using the ms-python.python.\r\n' +
-            `[Option 2] Set an interpreter using "${serverId}.interpreter" setting.\r\n` +
-            'Please use Python 3.8 or greater.',
+                '[Option 1] Select python interpreter using the ms-python.python.\r\n' +
+                `[Option 2] Set an interpreter using "${serverId}.interpreter" setting.\r\n` +
+                'Please use Python 3.8 or greater.',
         );
     };
 
@@ -125,7 +146,6 @@ export async function activate(context: vscode.ExtensionContext): Promise<void>
             if (result) {
                 statusBarItem.text = `$(kedro-logo)` + ' ' + result.label;
             }
-
         }),
         registerCommand('pygls.server.executeCommand', async () => {
             await executeServerCommand(lsClient);
@@ -152,4 +172,7 @@ export async function deactivate(): Promise<void> {
     if (lsClient) {
         await lsClient.stop();
     }
+    if (kedroVizProcess) {
+        kedroVizProcess.kill();
+    }
 }

src/webview/goToDefinition.ts

@@ -0,0 +1,49 @@
+import * as vscode from 'vscode';
+import { TextDocumentPositionParams } from 'vscode-languageclient';
+
+type Message = {
+    text: string;
+    type: string;
+};
+
+export async function goToDefinition(message: Message) {
+    const word = message.text;
+    let filePattern = '**/*.yml';
+
+    if (message.type === 'task') {
+        filePattern = '**/*.py';
+    }
+
+    const files = await vscode.workspace.findFiles(filePattern);
+
+    for (const file of files) {
+        const document = await vscode.workspace.openTextDocument(file);
+        const text = document.getText();
+        const regex = new RegExp(`\\b${word}\\b`, 'g');
+        let match;
+
+        while ((match = regex.exec(text)) !== null) {
+            const position = document.positionAt(match.index);
+            const params: TextDocumentPositionParams = {
+                textDocument: { uri: document.uri.toString() },
+                position,
+            };
+
+            const definitions = await vscode.commands.executeCommand<vscode.Location[]>(
+                'vscode.executeDefinitionProvider',
+                document.uri,
+                position,
+            );
+
+            if (definitions && definitions.length > 0) {
+                await vscode.window.showTextDocument(definitions[0].uri, {
+                    selection: definitions[0].range,
+                    viewColumn: vscode.ViewColumn.One,
+                });
+                return;
+            }
+        }
+    }
+
+    vscode.window.showInformationMessage(`No definition found for: ${word}`);
+}

src/webview/vizServer.ts

@@ -0,0 +1,36 @@
+import * as vscode from 'vscode';
+import { getWorkspaceFolders } from '../common/vscodeapi';
+import { fetchAndUpdateProjectData } from '../common/utilities';
+import { getInterpreterDetails } from '../common/python';
+const { exec } = require('child_process');
+
+async function getActivePythonInterpreter(): Promise<string | undefined> {
+    const interpreterDetails = await getInterpreterDetails();
+    if (interpreterDetails?.path) {
+        return interpreterDetails.path[0];
+    }
+    vscode.window.showErrorMessage('Python interpreter could not be resolved.');
+    return undefined;
+}
+
+export async function runKedroVizServer() {
+    const pythonPath = await getActivePythonInterpreter();
+    if (!pythonPath) return;
+
+    const workspaceFolders = getWorkspaceFolders();
+    if (!workspaceFolders.length) {
+        vscode.window.showErrorMessage('No workspace folder is open.');
+        return;
+    }
+    const workspacePath = workspaceFolders[0].uri.fsPath;
+
+    const command = `${pythonPath} -m kedro viz --no-browser -a --port=3131`;
+    const kedroVizProcess = exec(command, { cwd: workspacePath });
+
+    kedroVizProcess.stdout.on('data', (data: any) => {
+        console.log('Kedro Viz: ', data);
+        fetchAndUpdateProjectData();
+    });
+
+    return kedroVizProcess;
+}