browser: add screenshot functionality

src/overlay/obs/browser.py

@@ -1,9 +1,7 @@
-import sys
 import os
 import tempfile
 import shutil
 import base64
-import functools
 
 import trio
 from trio_cdp import page, emulation, dom
@@ -23,10 +21,10 @@ class BrowserCapturer:
         self.rate = rate
         self.bg_color = int(bg_color[1:3], 16), int(bg_color[3:5], 16), int(bg_color[5:7], 16), int(bg_color[7:9] or 'ff', 16)
 
-    def run(self):
-        trio.run(self.do_run, restrict_keyboard_interrupt_to_checkpoints=True)
+    def run_screencast(self):
+        trio.run(self.do_run_screencast, restrict_keyboard_interrupt_to_checkpoints=True)
 
-    async def do_run(self):
+    async def do_run_screencast(self):
         async with trio.open_nursery() as nursery:
             nursery.start_soon(self.capture_frames)
             nursery.start_soon(self.output_frames)
@@ -58,6 +56,30 @@ class BrowserCapturer:
                     await f.write(self.frame)
                 await trio.sleep_until(t + 1.0 / self.rate)
 
+    def run_screenshot(self):
+        trio.run(self.do_run_screenshot, restrict_keyboard_interrupt_to_checkpoints=True)
+
+    async def do_run_screenshot(self):
+        return await self.capture_shot()
+
+    async def capture_shot(self):
+        async with await self.chromium.connect() as conn, await self.chromium.select(conn) as sess, sess.page_enable():
+            # Set viewport size and transparency.
+            await emulation.set_device_metrics_override(
+                width=self.dimensions[0], height=self.dimensions[1],
+                device_scale_factor=1.0, mobile=False
+            )
+            await emulation.set_default_background_color_override(dom.RGBA(*self.bg_color))
+
+            # Navigate to the website.
+            async with sess.wait_for(page.LoadEventFired):
+                await page.navigate(self.url)
+
+            # Start capture.
+            frame = base64.b64decode(await page.capture_screenshot(format='png'))
+            async with await trio.open_file(self.outfile, 'wb') as f:
+                await f.write(frame)
+
 
 @OBSSource.type('browser_source')
 class OBSBrowserSource(OBSSource):
@@ -92,6 +114,6 @@ class OBSBrowserSource(OBSSource):
         runner = BrowserCapturer(self.url, fifo_path, dimensions=(self.width, self.height), bg_color=self.bg_color)
         return [FFmpegChain(
             FFmpegFilter('colorkey', color='0x' + self.bg_color[1:], similarity=0.3, blend=0.2),
-            inputs=[FFmpegInput(fifo_path, f='image2pipe', framerate=runner.rate, follow=1, probesize=8192)],
-            runners=[runner.run]
+            inputs=[FFmpegInput(fifo_path, f='image2pipe', r=runner.rate, follow=1, probesize=8192)],
+            runners=[runner.run_screencast]
         )]