Merge pull request #18 from IlIIIIIIlI/feature/export-flow-chart-as-png

Add PNG export functionality for flow charts

package.json

@@ -69,6 +69,7 @@
     "typescript": "^5.5.4"
   },
   "dependencies": {
+    "@microsoft/vscode-file-downloader-api": "^1.0.1",
     "@prisma/internals": "^5.19.1"
   }
 }

src/extension.ts

@@ -1,10 +1,13 @@
+import * as vscode from "vscode"
 import { getDMMF, getSchemaWithPath } from "@prisma/internals"
-import vscode from "vscode"
-
 import { transformDmmfToModelsAndConnections } from "./core/render"
 import { PrismaUMLPanel } from "./panels/prisma-uml-panel"
+let outputChannel: vscode.OutputChannel
 
 export function activate(context: vscode.ExtensionContext) {
+  outputChannel = vscode.window.createOutputChannel("Prisma Generate UML")
+  outputChannel.appendLine("Prisma Generate UML extension activated")
+
   const disposable = vscode.commands.registerCommand(
     "prisma-generate-uml.generateUML",
     async () => {
@@ -46,7 +49,13 @@ export function activate(context: vscode.ExtensionContext) {
         const { models, connections, enums } =
           transformDmmfToModelsAndConnections(response)
 
-        PrismaUMLPanel.render(context.extensionUri, models, connections, enums)
+        PrismaUMLPanel.render(
+          context.extensionUri,
+          models,
+          connections,
+          enums,
+          currentFileUri,
+        )
       } else {
         vscode.window.showInformationMessage(
           "Open a .prisma file to use this command",

src/panels/prisma-uml-panel.ts

@@ -1,23 +1,18 @@
-import {
-  Disposable,
-  Webview,
-  WebviewPanel,
-  window,
-  Uri,
-  ViewColumn,
-} from "vscode"
+import * as vscode from "vscode"
 import { getUri } from "../utilities/getUri"
 import { getNonce } from "../utilities/getNonce"
 import { Enum, Model, ModelConnection } from "../core/render"
 
 export class PrismaUMLPanel {
   public static currentPanel: PrismaUMLPanel | undefined
-  private readonly _panel: WebviewPanel
-  private _disposables: Disposable[] = []
+  public static readonly viewType = "prismaUML"
+  private readonly _panel: vscode.WebviewPanel
+  private _disposables: vscode.Disposable[] = []
 
   private constructor(
-    panel: WebviewPanel,
-    extensionUri: Uri,
+    panel: vscode.WebviewPanel,
+    private readonly _extensionUri: vscode.Uri,
+    private readonly _currentFileUri: vscode.Uri,
     models: Model[],
     connections: ModelConnection[],
     enums: Enum[],
@@ -26,12 +21,12 @@ export class PrismaUMLPanel {
 
     this._panel.onDidDispose(() => this.dispose(), null, this._disposables)
 
-    this._panel.webview.html = this._getWebviewContent(
-      this._panel.webview,
-      extensionUri,
-    )
+    this._panel.webview.html = this._getWebviewContent(this._panel.webview)
 
-    this._panel.iconPath = Uri.joinPath(extensionUri, "media/uml.svg")
+    this._panel.iconPath = vscode.Uri.joinPath(
+      this._extensionUri,
+      "media/uml.svg",
+    )
 
     this._panel.webview.postMessage({
       command: "setData",
@@ -42,41 +37,75 @@ export class PrismaUMLPanel {
 
     this._panel.webview.postMessage({
       command: "setTheme",
-      theme: window.activeColorTheme.kind,
+      theme: vscode.window.activeColorTheme.kind,
     })
+
+    this._panel.webview.onDidReceiveMessage(
+      async (message) => {
+        switch (message.command) {
+          case "saveImage":
+            await this._saveImage(message.data)
+            return
+        }
+      },
+      null,
+      this._disposables,
+    )
   }
 
   public static render(
-    extensionUri: Uri,
+    extensionUri: vscode.Uri,
     models: Model[],
     connections: ModelConnection[],
     enums: Enum[],
+    currentFileUri: vscode.Uri,
   ) {
     if (PrismaUMLPanel.currentPanel) {
-      PrismaUMLPanel.currentPanel._panel.reveal(ViewColumn.One)
+      PrismaUMLPanel.currentPanel._panel.reveal(vscode.ViewColumn.One)
     } else {
-      const panel = window.createWebviewPanel(
-        "prismaUML",
+      const panel = vscode.window.createWebviewPanel(
+        PrismaUMLPanel.viewType,
         "Prisma Schema UML",
-        ViewColumn.Two,
+        vscode.ViewColumn.Two,
         {
           enableScripts: true,
+          retainContextWhenHidden: true,
           localResourceRoots: [
-            Uri.joinPath(extensionUri, "out"),
-            Uri.joinPath(extensionUri, "webview-ui/build"),
+            vscode.Uri.joinPath(extensionUri, "out"),
+            vscode.Uri.joinPath(extensionUri, "webview-ui/build"),
           ],
         },
       )
       PrismaUMLPanel.currentPanel = new PrismaUMLPanel(
         panel,
         extensionUri,
+        currentFileUri,
         models,
         connections,
         enums,
       )
     }
   }
 
+  private async _saveImage(data: { format: string; dataUrl: string }) {
+    const base64Data = data.dataUrl.replace(/^data:image\/\w+;base64,/, "")
+    const buffer = Buffer.from(base64Data, "base64")
+
+    const uri = await vscode.window.showSaveDialog({
+      filters: { Images: [data.format] },
+      defaultUri: vscode.Uri.file(`prisma-uml.${data.format}`),
+    })
+
+    if (uri) {
+      try {
+        await vscode.workspace.fs.writeFile(uri, buffer)
+        vscode.window.showInformationMessage(`Image saved to ${uri.fsPath}`)
+      } catch (error) {
+        vscode.window.showErrorMessage(`Failed to save image: ${error}`)
+      }
+    }
+  }
+
   public dispose() {
     PrismaUMLPanel.currentPanel = undefined
     this._panel.dispose()
@@ -88,14 +117,14 @@ export class PrismaUMLPanel {
     }
   }
 
-  private _getWebviewContent(webview: Webview, extensionUri: Uri) {
-    const stylesUri = getUri(webview, extensionUri, [
+  private _getWebviewContent(webview: vscode.Webview) {
+    const stylesUri = getUri(webview, this._extensionUri, [
       "webview-ui",
       "build",
       "assets",
       "index.css",
     ])
-    const scriptUri = getUri(webview, extensionUri, [
+    const scriptUri = getUri(webview, this._extensionUri, [
       "webview-ui",
       "build",
       "assets",
@@ -109,7 +138,7 @@ export class PrismaUMLPanel {
         <head>
           <meta charset="UTF-8">
           <meta name="viewport" content="width=device-width, initial-scale=1.0">
-          <meta http-equiv="Content-Security-Policy" content="default-src 'none'; style-src ${webview.cspSource}; script-src 'nonce-${nonce}';">
+          <meta http-equiv="Content-Security-Policy" content="default-src 'none'; style-src ${webview.cspSource}; img-src ${webview.cspSource} data:; script-src 'nonce-${nonce}';">
           <link rel="stylesheet" type="text/css" href="${stylesUri}">
           <title>Prisma UML</title>
         </head>

webview-ui/package.json

@@ -11,6 +11,7 @@
   },
   "dependencies": {
     "@dagrejs/dagre": "^1.1.4",
+    "html-to-image": "^1.11.11",
     "react": "^18.3.1",
     "react-dom": "^18.3.1",
     "reactflow": "^11.11.4"

webview-ui/src/components/SchemaVisualizer.tsx

@@ -2,21 +2,25 @@ import ReactFlow, {
   Background,
   BackgroundVariant,
   ConnectionLineType,
+  ControlButton,
   Controls,
   MiniMap,
   Panel,
+  useReactFlow,
 } from "reactflow"
 import { useTheme } from "../lib/contexts/theme"
 import { useGraph } from "../lib/hooks/useGraph"
 import { Enum, Model, ModelConnection } from "../lib/types/schema"
-import { EnumNode } from "./EnumNode"
-import { ModelNode } from "./ModelNode"
 import {
   getButtonStyle,
   maskColor,
   nodeColor,
   nodeStrokeColor,
 } from "../lib/utils/colots"
+import { screenshot } from "../lib/utils/screnshot"
+import { EnumNode } from "./EnumNode"
+import { IDownload } from "./icons/IDownload"
+import { ModelNode } from "./ModelNode"
 
 interface Props {
   models: Model[]
@@ -26,6 +30,7 @@ interface Props {
 
 export const SchemaVisualizer = ({ connections, models, enums }: Props) => {
   const { isDarkMode } = useTheme()
+  const { getNodes } = useReactFlow()
 
   const modelNodes = models.map((model) => ({
     id: model.name,
@@ -62,7 +67,7 @@ export const SchemaVisualizer = ({ connections, models, enums }: Props) => {
 
   return (
     <div
-      className={`h-[100vh] w-full ${
+      className={`h-[100vh] w-full relative ${
         isDarkMode ? "bg-[#1c1c1c]" : "bg-[#e0e0e0]"
       }`}
     >
@@ -77,7 +82,11 @@ export const SchemaVisualizer = ({ connections, models, enums }: Props) => {
         minZoom={0.2}
         fitView
       >
-        <Controls />
+        <Controls>
+          <ControlButton title="Download" onClick={() => screenshot(getNodes)}>
+            <IDownload />
+          </ControlButton>
+        </Controls>
         <MiniMap
           nodeStrokeWidth={3}
           zoomable

webview-ui/src/components/icons/IDownload.tsx

@@ -0,0 +1,14 @@
+export const IDownload = () => (
+  <svg
+    width="16"
+    height="16"
+    viewBox="0 0 16 16"
+    fill="none"
+    xmlns="http://www.w3.org/2000/svg"
+  >
+    <path
+      d="M8 12L3 7L4.4 5.55L7 8.15V0H9V8.15L11.6 5.55L13 7L8 12ZM2 16C1.45 16 0.979333 15.8043 0.588 15.413C0.196666 15.0217 0.000666667 14.5507 0 14V11H2V14H14V11H16V14C16 14.55 15.8043 15.021 15.413 15.413C15.0217 15.805 14.5507 16.0007 14 16H2Z"
+      fill="black"
+    />
+  </svg>
+)

webview-ui/src/lib/utils/screnshot.ts

@@ -0,0 +1,58 @@
+import { toPng } from "html-to-image"
+import { getNodesBounds, getViewportForBounds, Node } from "reactflow"
+
+interface VSCodeAPI {
+  postMessage(message: SaveImageMessage): void
+  getState(): unknown
+  setState(state: unknown): void
+}
+
+interface SaveImageMessage {
+  command: "saveImage"
+  data: {
+    format: "png"
+    dataUrl: string
+  }
+}
+
+declare function acquireVsCodeApi(): VSCodeAPI
+
+const vscode = acquireVsCodeApi()
+
+export const screenshot = (getNodes: () => Node[]) => {
+  const nodesBounds = getNodesBounds(getNodes())
+
+  // 8k resolution
+  const imageWidth = 7680
+  const imageHeight = 4320
+
+  const transform = getViewportForBounds(
+    nodesBounds,
+    imageWidth,
+    imageHeight,
+    0,
+    2,
+  )
+
+  toPng(document.querySelector(".react-flow__viewport") as HTMLElement, {
+    filter: (node) => {
+      const exclude = ["react-flow__minimap", "react-flow__controls"]
+      return !exclude.some((className) => node.classList?.contains(className))
+    },
+    backgroundColor: "transparent",
+    width: imageWidth,
+    height: imageHeight,
+    style: {
+      transform: `translate(${transform.x}px, ${transform.y}px) scale(${transform.zoom})`,
+    },
+  })
+    .then((dataUrl) => {
+      vscode.postMessage({
+        command: "saveImage",
+        data: { format: "png", dataUrl },
+      })
+    })
+    .catch((error) => {
+      console.error("Error generating image:", error)
+    })
+}

webview-ui/tsconfig.app.tsbuildinfo

@@ -1 +1 @@
-{"root":["./src/app.tsx","./src/index.tsx","./src/vite-env.d.ts","./src/components/enumnode.tsx","./src/components/modelnode.tsx","./src/components/schemavisualizer.tsx","./src/lib/contexts/theme.tsx","./src/lib/hooks/usegraph.ts","./src/lib/types/schema.ts","./src/lib/utils/colots.ts","./src/lib/utils/layout-utils.ts"],"version":"5.6.2"}
+{"root":["./src/app.tsx","./src/index.tsx","./src/vite-env.d.ts","./src/components/enumnode.tsx","./src/components/modelnode.tsx","./src/components/schemavisualizer.tsx","./src/components/icons/idownload.tsx","./src/lib/contexts/theme.tsx","./src/lib/hooks/usegraph.ts","./src/lib/types/schema.ts","./src/lib/utils/colots.ts","./src/lib/utils/layout-utils.ts","./src/lib/utils/screnshot.ts"],"version":"5.6.2"}