[REPL] Add Python execution support

package.json

@@ -462,9 +462,14 @@
       },
       {
         "category": "KX",
-        "command": "kdb.start.repl",
+        "command": "kdb.repl.start",
         "title": "Start REPL"
       },
+      {
+        "category": "KX",
+        "command": "kdb.repl.restart",
+        "title": "Restart REPL"
+      },
       {
         "category": "KX",
         "command": "kdb.execute.selectedQuery",
@@ -1081,6 +1086,12 @@
           "group": "q",
           "when": "resourceExtname == .py"
         }
+      ],
+      "terminal/context": [
+        {
+          "command": "kdb.repl.restart",
+          "group": "q"
+        }
       ]
     },
     "customEditors": [
@@ -1194,6 +1205,7 @@
     "prettier": "^3.6.0",
     "rimraf": "^6.0.1",
     "sinon": "^21.0.0",
+    "tree-kill": "^1.2.2",
     "typescript": "^5.8.3",
     "typescript-eslint": "^8.34.1",
     "vscode-extension-tester": "^8.16.2",

src/classes/replConnection.ts

@@ -12,6 +12,8 @@
  */
 
 import { ChildProcessWithoutNullStreams, spawn } from "node:child_process";
+import path from "node:path";
+import kill from "tree-kill";
 import * as vscode from "vscode";
 
 import { ext } from "../extensionVariables";
@@ -41,6 +43,7 @@ const ANSI = {
 const KEY = {
   CR: "\r",
   CTRLC: "\x03",
+  CTRLD: "\x04",
   BS: "\b",
   BSMAC: "\x7f",
   DEL: "\x1b[3~",
@@ -95,6 +98,7 @@ export interface Result {
 }
 
 export class ReplConnection {
+  private readonly posix = process.platform !== "win32";
   private readonly history = new History();
   private readonly identity = crypto.randomUUID();
   private readonly token = new RegExp(
@@ -128,6 +132,8 @@ export class ReplConnection {
   private exited = false;
   private stopped = false;
   private executing?: Execution;
+  private activate = "";
+  private script = "";
 
   private constructor() {
     this.onDidWrite = new vscode.EventEmitter<string>();
@@ -184,8 +190,17 @@ export class ReplConnection {
   }
 
   private createProcess() {
-    return spawn(getQExecutablePath(), ["-q"], {
+    const params = ["-q"];
+
+    if (this.script) {
+      params.unshift(this.script);
+      this.script = "";
+    }
+    const q = getQExecutablePath();
+
+    return spawn(this.activate ? `${this.activate};${q}` : q, params, {
       env: { ...process.env, QHOME: ext.REAL_QHOME },
+      shell: this.activate ? "bash" : false,
     });
   }
 
@@ -235,12 +250,19 @@ export class ReplConnection {
   }
 
   private stopExecution() {
-    this.stopped = process.platform === "win32";
-    this.process.kill("SIGINT");
+    this.stopped = !this.posix;
+    const signal = "SIGINT";
+    if (this.activate) {
+      if (this.process.pid) kill(this.process.pid, signal);
+    } else this.process.kill(signal);
   }
 
-  private stopProcess() {
-    this.process.kill("SIGTERM");
+  private stopProcess(restart = false) {
+    this.stopped = restart;
+    const signal = restart ? "SIGKILL" : "SIGTERM";
+    if (this.activate) {
+      if (this.process.pid) kill(this.process.pid, signal);
+    } else this.process.kill(signal);
   }
 
   private runQuery(data: string) {
@@ -347,6 +369,11 @@ export class ReplConnection {
     }
   }
 
+  private restart() {
+    this.cancel();
+    this.stopProcess(true);
+  }
+
   private executeNext() {
     if (!this.executing && !this.messages) {
       this.executing = this.executions.shift();
@@ -357,14 +384,20 @@ export class ReplConnection {
     }
   }
 
+  private normalize(decoded: string) {
+    return decoded.replace(/(?:\r\n|[\r\n])+/gs, ANSI.CRLF);
+  }
+
+  private clean(decoded: string) {
+    return decoded.replace(/(?:\r\n|[\r\n])+/gs, "");
+  }
+
   private push(data: any, buffer: string[]) {
     const c = this.executing;
     if (!c) return;
     const decoded = this.decoder.decode(data);
     this.token.lastIndex = 0;
-    const output = decoded
-      .replace(this.token, ANSI.EMPTY)
-      .replace(/(?:\r\n|[\r\n])+/gs, ANSI.CRLF);
+    const output = this.normalize(decoded.replace(this.token, ANSI.EMPTY));
     if (output) {
       buffer.push(output);
       this.sendToTerminal(output);
@@ -405,9 +438,12 @@ export class ReplConnection {
 
   private handleOutput(data: any) {
     const c = this.executing;
-    if (!c) return;
-    const token = this.push(data, c.line);
-    if (token) this.nextToken(token);
+    if (c) {
+      const token = this.push(data, c.line);
+      if (token) this.nextToken(token);
+    } else {
+      this.sendToTerminal(this.normalize(this.decoder.decode(data)));
+    }
   }
 
   private handleError(error: Error) {
@@ -423,6 +459,8 @@ export class ReplConnection {
       this.resolve();
       this.process = this.createProcess();
       this.connect();
+      this.sendToTerminal(ANSI.CRLF);
+      this.showPrompt(true);
       return;
     }
     this.sendToTerminal(
@@ -441,7 +479,9 @@ export class ReplConnection {
   }
 
   private close() {
-    if (ReplConnection.instance === this) ReplConnection.instance = undefined;
+    if (ReplConnection.instance === this) {
+      ReplConnection.instance = undefined;
+    }
     this.cancel();
     this.stopProcess();
     this.onDidWrite.dispose();
@@ -460,25 +500,23 @@ export class ReplConnection {
       return;
     }
 
-    const inputText = this.inputText;
-
-    if (data === KEY.CR) {
-      if (this.executing) {
-        this.sendToTerminal(ANSI.CRLF);
-        return;
-      }
-      if (/^\\[\t ]*$/m.test(inputText)) {
-        this.context = this.context === CTX.K ? CTX.Q : CTX.K;
-        this.sendCommand("\\");
-        this.sendToTerminal(ANSI.CRLF);
-        this.inputText = ANSI.EMPTY;
-        this.showPrompt(true);
-        return;
-      }
-    }
+    let inputText: string | undefined;
 
     switch (data) {
       case KEY.CR:
+        if (this.executing) {
+          this.sendToTerminal(ANSI.CRLF);
+          break;
+        }
+        inputText = this.inputText;
+        if (/^\\[\t ]*$/m.test(inputText)) {
+          this.context = this.context === CTX.K ? CTX.Q : CTX.K;
+          this.sendCommand("\\");
+          this.sendToTerminal(ANSI.CRLF);
+          this.inputText = ANSI.EMPTY;
+          this.showPrompt(true);
+          break;
+        }
         this.history.push(inputText);
         this.history.rewind();
         this.runQuery(inputText);
@@ -490,6 +528,9 @@ export class ReplConnection {
       case KEY.CTRLC:
         this.cancel();
         break;
+      case KEY.CTRLD:
+        this.restart();
+        break;
       case KEY.BS:
       case KEY.BSMAC:
         if (this.inputIndex > 0 && this.input.splice(this.inputIndex - 1, 1)) {
@@ -550,7 +591,17 @@ export class ReplConnection {
         break;
       default:
         if (/(?:\r\n|[\r\n])/s.test(data)) {
-          if (!/\.venv[/\\]+bin[/\\]+activate/is.test(data)) {
+          if (/\.venv[/\\]+(?:scripts|bin)[/\\]+/is.test(data)) {
+            if (this.posix) {
+              this.activate = this.clean(data);
+              this.stopProcess(true);
+              this.sendToTerminal(ANSI.CRLF + "Restarting REPL on " + data);
+            }
+          } else if (path.isAbsolute(data)) {
+            this.script = this.clean(data);
+            this.stopProcess(true);
+            this.sendToTerminal(ANSI.CRLF + "Running " + this.script);
+          } else {
             this.runQuery(data);
           }
           break;
@@ -613,6 +664,12 @@ export class ReplConnection {
 
   private static instance?: ReplConnection;
 
+  static restart() {
+    if (this.instance && !this.instance.exited) {
+      this.instance.restart();
+    }
+  }
+
   static getOrCreateInstance() {
     if (!this.instance || this.instance.exited) {
       this.instance = new ReplConnection();

src/commands/workspaceCommand.ts

@@ -46,6 +46,7 @@ import {
   notify,
   Runner,
 } from "../utils/notifications";
+import { getPythonWrapper } from "../utils/queryUtils";
 import {
   cleanAssemblyName,
   cleanDapName,
@@ -219,7 +220,7 @@ export async function pickConnection(uri: Uri) {
   const servers = getServers();
 
   const items = ["(none)"];
-  if (isQ(uri) || isNotebook(uri)) {
+  if (isQ(uri) || isNotebook(uri) || isPython(uri)) {
     items.push(ext.REPL);
   }
   items.push(...servers);
@@ -462,9 +463,10 @@ function isKxFolder(uri: Uri | undefined) {
   return uri && Path.basename(uri.path) === ".kx";
 }
 
-export async function startRepl() {
+export async function startRepl(restart = false) {
   try {
-    ReplConnection.getOrCreateInstance().start();
+    if (restart) ReplConnection.restart();
+    else ReplConnection.getOrCreateInstance().start();
   } catch (error) {
     notify(errorMessage(error), MessageKind.ERROR, {
       logger,
@@ -474,31 +476,39 @@ export async function startRepl() {
 }
 
 export async function runOnRepl(editor: TextEditor, type?: ExecutionTypes) {
-  const basename = getBasename(editor.document.uri);
+  const uri = editor.document.uri;
+  const basename = getBasename(uri);
 
   let text: string;
 
-  if (type === ExecutionTypes.QueryFile) {
-    text = editor.document.getText();
-  } else if (type === ExecutionTypes.QuerySelection) {
-    const selection = editor.selection;
-    text = selection.isEmpty
-      ? editor.document.lineAt(selection.active.line).text
-      : editor.document.getText(selection);
-  } else {
-    notify(
-      `Executing ${basename} on ${ext.REPL} is not supported.`,
-      MessageKind.ERROR,
-      { logger },
-    );
-    return;
+  switch (type) {
+    case ExecutionTypes.QueryFile:
+    case ExecutionTypes.PythonQueryFile:
+      text = editor.document.getText();
+      break;
+    case ExecutionTypes.QuerySelection:
+    case ExecutionTypes.PythonQuerySelection:
+      text = editor.selection.isEmpty
+        ? editor.document.lineAt(editor.selection.active.line).text
+        : editor.document.getText(editor.selection);
+      break;
+    default:
+      notify(
+        `Executing ${basename} on ${ext.REPL} is not supported.`,
+        MessageKind.ERROR,
+        { logger },
+      );
+      return;
   }
 
   try {
     const runner = Runner.create((_, token) => {
       const repl = ReplConnection.getOrCreateInstance();
       repl.show();
-      return repl.executeQuery(text, token);
+      return repl.executeQuery(
+        isPython(uri) ? getPythonWrapper(text) : text,
+        token,
+      );
     });
     runner.cancellable = Cancellable.EXECUTOR;
     runner.title = `Executing ${basename} on ${ext.REPL}.`;