From a49be3aa4d827993e922917af0dff6ce54438d47 Mon Sep 17 00:00:00 2001
From: Brent Yi <yibrenth@gmail.com>
Date: Thu, 26 Oct 2023 13:52:18 -0700
Subject: [PATCH] Add `send_file_download()` (#119)

* Add `send_file_download()`

* Add link remove

---------

Co-authored-by: Rohan Mathur <jmlife6@gmail.com>
---
 src/viser/_message_api.py                   | 15 +++++++++++++++
 src/viser/_messages.py                      |  9 +++++++++
 src/viser/client/src/WebsocketInterface.tsx | 11 ++++++++++-
 src/viser/client/src/WebsocketMessages.tsx  | 13 ++++++++++++-
 4 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/src/viser/_message_api.py b/src/viser/_message_api.py
index 148808e6f..8725851ac 100644
--- a/src/viser/_message_api.py
+++ b/src/viser/_message_api.py
@@ -11,6 +11,7 @@
 import base64
 import colorsys
 import io
+import mimetypes
 import queue
 import threading
 import time
@@ -819,3 +820,17 @@ def add_3d_gui_container(
         )
         node_handle = SceneNodeHandle._make(self, name, wxyz, position, visible=visible)
         return Gui3dContainerHandle(node_handle._impl, gui_api, container_id)
+
+    def send_file_download(self, filename: str, content: bytes) -> None:
+        """Send a file for a client or clients to download."""
+        mime_type = mimetypes.guess_type(filename, strict=False)[0]
+        assert (
+            mime_type is not None
+        ), f"Could not guess MIME type from filename {filename}!"
+        self._queue(
+            _messages.FileDownload(
+                filename,
+                content=content,
+                mime_type=mime_type,
+            )
+        )
diff --git a/src/viser/_messages.py b/src/viser/_messages.py
index 7971fb542..c4ffbcf8f 100644
--- a/src/viser/_messages.py
+++ b/src/viser/_messages.py
@@ -565,3 +565,12 @@ class GetRenderResponseMessage(Message):
     """Message from client->server carrying a render."""
 
     payload: bytes
+
+
+@dataclasses.dataclass
+class FileDownload(Message):
+    """Send a file for clients to download."""
+
+    filename: str
+    content: bytes
+    mime_type: str
diff --git a/src/viser/client/src/WebsocketInterface.tsx b/src/viser/client/src/WebsocketInterface.tsx
index 652849d98..5469311d7 100644
--- a/src/viser/client/src/WebsocketInterface.tsx
+++ b/src/viser/client/src/WebsocketInterface.tsx
@@ -141,7 +141,7 @@ function useMessageHandler() {
                 rotation={
                   // There's redundancy here when we set the side to
                   // THREE.DoubleSide, where xy and yx should be the same.
-                  // 
+                  //
                   // But it makes sense to keep this parameterization because
                   // specifying planes by xy seems more natural than the normal
                   // direction (z, +z, or -z), and it opens the possibility of
@@ -734,6 +734,15 @@ function useMessageHandler() {
         );
         return;
       }
+      case "FileDownload": {
+        const blob = new Blob([message.content], { type: message.mime_type });
+        const link = document.createElement("a");
+        link.href = window.URL.createObjectURL(blob);
+        link.download = message.filename;
+        link.click();
+        link.remove();
+        return;
+      }
       default: {
         console.log("Received message did not match any known types:", message);
         return;
diff --git a/src/viser/client/src/WebsocketMessages.tsx b/src/viser/client/src/WebsocketMessages.tsx
index ab6f6708e..8e8d934cc 100644
--- a/src/viser/client/src/WebsocketMessages.tsx
+++ b/src/viser/client/src/WebsocketMessages.tsx
@@ -669,6 +669,16 @@ export interface GetRenderResponseMessage {
   type: "GetRenderResponseMessage";
   payload: Uint8Array;
 }
+/** Send a file for clients to download.
+ *
+ * (automatically generated)
+ */
+export interface FileDownload {
+  type: "FileDownload";
+  filename: string;
+  content: Uint8Array;
+  mime_type: string;
+}
 
 export type Message =
   | ViewerCameraMessage
@@ -723,4 +733,5 @@ export type Message =
   | CatmullRomSplineMessage
   | CubicBezierSplineMessage
   | GetRenderRequestMessage
-  | GetRenderResponseMessage;
+  | GetRenderResponseMessage
+  | FileDownload;