save_screenshot seems to be broken on standalone server

[fischman]

Is save_screenshot meant to work in the standalone server?
AFAICT it can't work with standalone.py as written due to mismatched types/expectations, but even if those are fixed it seems to rely on the browser being able to write to the file-system, which seems impossible without some sort of permissions flow or download file dialog / picker(?), or a return message from JS to the standalone server, but I'm not seeing evidence that that is happening, either.

Reproduction recipe:

1. Launch: python -m ocp_vscode and open http://localhost:3939/viewer in a browser
2. In another terminal: python -c 'from ocp_vscode import show, save_screenshot; from build123d import Box; show(Box(1,2,3)); print(save_screenshot("x.png"))'

Expected:
a) The box is rendered in the browser's viewer
b) A screenshot of the box is written to a file named "x.png" in the standalone server's CWD.

Actual:
i) The box is rendered in the browser's viewer as expected (yay)
ii) save_screenshot raises an exception; when a naive fix is applied to standalone.py then it takes 2s, returns None, and no file is written out. Using an absolute path like /tmp/x.png instead of the relative one above doesn't seem to matter.

---

Naive fix mentioned above is:

~/tmp/vscode-ocp-cad-viewer/ocp_vscode {main} $ git diff
diff --git a/ocp_vscode/standalone.py b/ocp_vscode/standalone.py
index 89a15f4..f3cd1c7 100644
--- a/ocp_vscode/standalone.py
+++ b/ocp_vscode/standalone.py
@@ -210,9 +210,9 @@ class Viewer:
                     self.debug_print("Received config command")
                     self.config["_splash"] = self.splash
                     self.python_client.send(orjson.dumps(self.config))
-                elif cmd.type == "screenshot":
+                elif cmd.get("type") == "screenshot":
                     self.debug_print("Received screenshot command")
-                    self.python_client(orjson.dumps(cmd))
+                    self.python_client.send(orjson.dumps(cmd))

             elif message_type == "D":
                 self.python_client = ws

[bernhard-42]

Well, I guess I never tested this, sorry. I'll look into it

[bernhard-42]

yep, not only untested, but pretty much not implemented :-)
Seems I have forgotten to add it to my TODO list these days.
It works now and will be in the next release.

This is how it works:

• Python sends screenshot command with the the png save-path to Javasscript via websockets
• Viewer creates screenshot
• Viewer sends png data as data url together with the received filename back to Python over websockets
• Python stores png data under the received filename.

So, browser does not need any dangerous "write-to-local-filesystem"

Caveat: The process is completely async (well, we talk to Javascript at the end), and by passing the filename through it'll work. But, you might not be able to rely on the screenshot being written immediately after save_screenshot returns in Python. Only when the websocket command comes back, Python will write the file

[fischman]

Thanks, @bernhard-42 !
Can you share the fix commit so I can patch it locally until the next release? (Or do you expect the next release to be imminent?)

Polling client side for appearance of the saved file seems fine to me. I hope that the sync aspects of the implementation are such that as soon as the save_screenshot call returns, it is at least guaranteed that the snapshot has been taken (in js), so that it is safe to call other Python show() commands and know that the screenshot will not reflect these later show()s. If true, and if your fix commit removes the polling in save_screenshot, then this would enable a pipelined approach to dumping screenshots from a successive set of views without having to wait for each file to be written to request the next screenshot.

[bernhard-42]

This is the fix

diff --git a/ocp_vscode/standalone.py b/ocp_vscode/standalone.py
index 29682ca..70491b1 100644
--- a/ocp_vscode/standalone.py
+++ b/ocp_vscode/standalone.py
@@ -1,3 +1,4 @@
+import base64
 import orjson
 import socket
 import yaml
@@ -93,6 +94,17 @@ def COMMS(host, port):
 INIT = """onload="showViewer()" """


+def save_png_data_url(data_url, output_path):
+    base64_data = data_url.split(",")[1]
+    image_data = base64.b64decode(base64_data)
+    try:
+        with open(output_path, "wb") as f:
+            f.write(image_data)
+        print(f"Write png file to {output_path}")
+    except Exception as ex:
+        print("Cannot save png file:", str(ex))
+
+
 class Viewer:
     def __init__(self, params):
         self.status = {}
@@ -224,14 +236,19 @@ class Viewer:
                 if cmd == "status":
                     self.debug_print("Received status command")
                     self.python_client.send(orjson.dumps({"text": self.status}))
+
                 elif cmd == "config":
                     self.debug_print("Received config command")
                     self.configure(self.params)
                     self.config["_splash"] = self.splash
                     self.python_client.send(orjson.dumps(self.config))
-                elif cmd.type == "screenshot":
+
+                elif cmd.get("type") == "screenshot":
                     self.debug_print("Received screenshot command")
-                    self.python_client(orjson.dumps(cmd))
+                    if self.javascript_client is None:
+                        self.not_registered()
+                        continue
+                    self.javascript_client.send(data)

             elif message_type == "D":
                 self.python_client = ws
@@ -245,11 +262,18 @@ class Viewer:

             elif message_type == "U":
                 self.javascript_client = ws
-                changes = orjson.loads(data)["text"]
-                self.debug_print("Received incremental UI changes", changes)
-                for key, value in changes.items():
-                    self.status[key] = value
-                self.backend.handle_event(changes, MessageType.UPDATES)
+                message = orjson.loads(data)
+                if message["command"] == "screenshot":
+                    filename = message["text"]["filename"]
+                    data_url = message["text"]["data"]
+                    self.debug_print("Received screenshot data for file", filename)
+                    save_png_data_url(data_url, filename)
+                else:
+                    changes = message["text"]
+                    self.debug_print("Received incremental UI changes", changes)
+                    for key, value in changes.items():
+                        self.status[key] = value
+                    self.backend.handle_event(changes, MessageType.UPDATES)

             elif message_type == "S":
                 self.python_client = ws

[bernhard-42]

I still need to do more testing to check whether there are regressions.
But from a pure screenshot perspective it worked for me

[fischman]

Thanks, that seems like it works to me too.

Suggestions:

1. Make the new save_png_data_url write to a tempfile (in the same directory as the requested output file with a unique suffix, for example) and rename it into the requested filename after done writing, to ensure that a polling client doesn't read a partially-written file (b/c os.rename is atomic at least on POSIX inside modern filesystems).
2. Drop the polling in show.py:save_screenshot. Enforcing polling in ocp_vscode prevents the client from being able to pipeline paired show+save_screenshot calls and only poll after all are requested. If you do want to keep the polling where it is, some observations on opportunities for improvement:
   1. it's almost guaranteed to sleep at least 100ms regardless of how quickly the screenshot is written
   2. if it did poll for an arbitrary period (2s right now), it should probably indicate timeout in some way the caller can discern (return value or exception). Right now execution simply falls off the def (perhaps a return or raise is missing before/around that last string literal?)

[bernhard-42]

I'll add a flag to save_screenshot:

def save_screenshot(filename, port=None, polling=True):
    """Save a screenshot of the current view"""
    if not filename.startswith(os.sep):
        prefix = pathlib.Path(".").absolute()
        full_path = str(prefix / filename)
    else:
        full_path = filename
    p = pathlib.Path(full_path)
    mtime = p.stat().st_mtime if p.exists() else 0

    send_command({"type": "screenshot", "filename": f"{full_path}"}, port=port)

    if polling:
        done = False
        for i in range(20):
            if p.exists() and p.stat().st_mtime > mtime:
                print("Screenshot saved to ", full_path)
                done = True
                break
            time.sleep(0.1)

        if not done:
            print("Warning: Screenshot not found in 2 seconds, aborting")

So, for not-so-experts it'll wait. And if you want to avoid polling for your workflow, then add the parameter