Merge pull request #509 from KxSystems/ee-api

[ggplot] Move png detection to extension

Author: ecmel

README.md

@@ -629,7 +629,7 @@ To use GGPlot2 in VSCode:
 
 You can make changes to the script before exporting the plot. Re-running the script updates to reflect the changes.
 
-**Note**: When executing GG script commands in PyKX, select the `KDB RESULTS` tab to display the plot.
+**Note**: When executing GG script commands, select the `KDB RESULTS` tab to display the plot.
 
 ## q REPL
 

resources/evaluate.q

@@ -1,4 +1,8 @@
 {[ctx; code; returnFormat]
+  if [`histogram in key `.qp;
+  if [not `display2 in key `.qp; 
+  .qp.display2: (')[{x[`output][`bytes]}; .qp.display]
+  ]]; 
   if [-10h ~ type ctx;
   ctx: enlist ctx];
   toString: {[data]
@@ -117,8 +121,7 @@
   (`result;    ::);
   (`errored;   1b);
   (`error;     err);
-  (`backtrace; .Q.sbt userCode);
-  (`base64;    0b))
+  (`backtrace; .Q.sbt userCode))
   }[suffix; prefix]];
   if [isLastLine or result`errored;
   system "d ", cachedCtx;
@@ -197,23 +200,6 @@
   }
   result: evalInContext[ctx; splitExpression stripTrailingSemi wrapLines removeMultilineComments code];
   if [result `errored; :result];
-  /ggplot - start
-  if[type[result[`result]] = 99h;
-  attrs: key[[result[`result]]]; 
-  if[type[attrs] = 11h;
-  if[`output in attrs; 
-  output: result[`result][`output];
-  if[type[output] = 99h;
-  attrs: key[output];
-  if[type[attrs] = 11h;
-  if[`bytes in attrs;
-  bytes: output[`bytes];
-  if[type[bytes] = 4h;
-  if[0x89504E470D0A1A0A ~ bytes til 8;
-  result[`base64]: 1b; 
-  result[`result]: .Q.btoa bytes; 
-  :result]]]]]]]];
-  /ggplot - end
   if [returnFormat ~ "text";
   result[`result]: toString result `result];
   if [returnFormat ~ "structuredText";

resources/evaluatePy.q

@@ -1,4 +1,8 @@
 {[isString;code]
+ if [`histogram in key `.qp;
+ if [not `display2 in key `.qp; 
+ .qp.display2: (')[{x[`output][`bytes]}; .qp.display]
+ ]]; 
  if[()~key`.pykx;
  :(!). flip(
  (`result;::);
@@ -97,23 +101,6 @@
  if [`backtrace in key result;
  result[`backtrace]:string result`backtrace];
  if[result `errored; :result];
- /ggplot - start
- if[type[result[`result]] = 99h;
- attrs: key[[result[`result]]]; 
- if[type[attrs] = 11h;
- if[`output in attrs; 
- output: result[`result][`output];
- if[type[output] = 99h;
- attrs: key[output];
- if[type[attrs] = 11h;
- if[`bytes in attrs;
- bytes: output[`bytes];
- if[type[bytes] = 4h;
- if[0x89504E470D0A1A0A ~ bytes til 8;
- result[`base64]: 1b; 
- result[`result]: .Q.btoa bytes; 
- :result]]]]]]]];
- /ggplot - end
  $[result`errored;
  ::;
  returnResult;

src/classes/localConnection.ts

@@ -175,7 +175,6 @@ export class LocalConnection {
     isPython?: boolean,
   ): Promise<any> {
     let result;
-    let base64 = false;
     await this.waitForConnection();
 
     if (!this.connection) {
@@ -204,7 +203,6 @@ export class LocalConnection {
           );
         } else {
           result = res.result === null ? "" : res.result;
-          base64 = res.base64 || false;
         }
       }
     });
@@ -222,10 +220,6 @@ export class LocalConnection {
       return result;
     }
 
-    if (base64) {
-      return { base64, result };
-    }
-
     if (!stringify && !isPython) {
       return JSON.parse(result);
     }

src/commands/serverCommand.ts

@@ -58,6 +58,7 @@ import {
   checkIfIsDatasource,
   addQueryHistory,
   formatScratchpadStacktrace,
+  resultToBase64,
 } from "../utils/queryUtils";
 import {
   validateServerAlias,
@@ -894,37 +895,43 @@ async function _executeQuery(
     );
   } else {
     /* istanbul ignore next */
-    if (results.base64) {
-      const active = ext.activeTextEditor;
-      if (active) {
-        const data = `data:image/png;base64,${results.result}`;
-        const plot = <Plot>{
-          charts: [{ data }],
-        };
-        const uri = await addWorkspaceFile(
-          active.document.uri,
-          "plot",
-          ".plot",
-        );
-        if (!workspaceHas(uri)) {
-          await workspace.openTextDocument(uri);
-          await openWith(uri, ChartEditorProvider.viewType, ViewColumn.Beside);
+    if (ext.isResultsTabVisible) {
+      const data = resultToBase64(results);
+      if (data) {
+        const active = ext.activeTextEditor;
+        if (active) {
+          const plot = <Plot>{
+            charts: [{ data }],
+          };
+          const uri = await addWorkspaceFile(
+            active.document.uri,
+            "plot",
+            ".plot",
+          );
+          if (!workspaceHas(uri)) {
+            await workspace.openTextDocument(uri);
+            await openWith(
+              uri,
+              ChartEditorProvider.viewType,
+              ViewColumn.Beside,
+            );
+          }
+          await setUriContent(uri, JSON.stringify(plot));
         }
-        await setUriContent(uri, JSON.stringify(plot));
+      } else {
+        await writeQueryResultsToView(
+          results,
+          query,
+          connLabel,
+          executorName,
+          isInsights,
+          isWorkbook ? "WORKBOOK" : "SCRATCHPAD",
+          isPython,
+          duration,
+          isFromConnTree,
+          connVersion,
+        );
       }
-    } else if (ext.isResultsTabVisible) {
-      await writeQueryResultsToView(
-        results,
-        query,
-        connLabel,
-        executorName,
-        isInsights,
-        isWorkbook ? "WORKBOOK" : "SCRATCHPAD",
-        isPython,
-        duration,
-        isFromConnTree,
-        connVersion,
-      );
     } else {
       await writeQueryResultsToConsole(
         results,

src/models/queryResult.ts

@@ -26,7 +26,6 @@ export type QueryResult = {
     text: string;
     index: number;
   }[];
-  base64?: boolean;
 };
 
 export enum QueryResultType {

src/utils/queryUtils.ts

@@ -395,3 +395,19 @@ export function formatScratchpadStacktrace(stacktrace: ScratchpadStacktrace) {
     })
     .join("\n");
 }
+
+const PNG = ["0x89", "0x50", "0x4e", "0x47", "0x0d", "0x0a", "0x1a", "0x0a"];
+
+export function resultToBase64(result: any): string | undefined {
+  const bytes =
+    result?.columns?.type === "bytes" ? result.columns.values : result;
+  if (Array.isArray(bytes) && bytes.length > 66) {
+    for (let i = 0; i < PNG.length; i++) {
+      if (bytes[i] !== PNG[i] && bytes[i] !== parseInt(PNG[i], 16)) {
+        return undefined;
+      }
+    }
+    return `data:image/png;base64,${Buffer.from(bytes).toString("base64")}`;
+  }
+  return undefined;
+}

test/suite/utils.test.ts

@@ -1974,4 +1974,58 @@ describe("Utils", () => {
       );
     });
   });
+
+  describe("resultToBase64", () => {
+    const png = [
+      "0x89",
+      "0x50",
+      "0x4e",
+      "0x47",
+      "0x0d",
+      "0x0a",
+      "0x1a",
+      "0x0a",
+    ];
+    const img = Array.from({ length: 59 }, () => "0x00");
+
+    it("should return undefined for undefined", () => {
+      const result = queryUtils.resultToBase64(undefined);
+      assert.strictEqual(result, undefined);
+    });
+    it("should return undefined for just signature", () => {
+      const result = queryUtils.resultToBase64(png);
+      assert.strictEqual(result, undefined);
+    });
+    it("should return undefined for undefined for bad signature", () => {
+      const result = queryUtils.resultToBase64([
+        ...png.map((v) => parseInt(v, 16) + 1),
+        ,
+        ...img,
+      ]);
+      assert.strictEqual(result, undefined);
+    });
+    it("should return base64 for minimum img str", () => {
+      const result = queryUtils.resultToBase64([...png, ...img]);
+      assert.ok(result.startsWith("data:image/png;base64"));
+    });
+    it("should return base64 for minimum img num", () => {
+      const result = queryUtils.resultToBase64([
+        ...png.map((v) => parseInt(v, 16)),
+        ...img.map((v) => parseInt(v, 16)),
+      ]);
+      assert.ok(result.startsWith("data:image/png;base64"));
+    });
+    it("should return base64 for minimum img str for structuredText", () => {
+      const result = queryUtils.resultToBase64({
+        columns: { type: "bytes", values: [...png, ...img] },
+      });
+      assert.ok(result.startsWith("data:image/png;base64"));
+    });
+    it("should return undefined for bogus structuredText", () => {
+      const result = queryUtils.resultToBase64({
+        columns: { type: "bytes" },
+      });
+      assert.strictEqual(result, undefined);
+    });
+  });
 });