From ab990dba1405823583c80c0182f6536b5eca065e Mon Sep 17 00:00:00 2001
From: "Dustin C. Hatch" <dustin@hatch.name>
Date: Sat, 30 Apr 2022 15:02:32 -0500
Subject: [PATCH] Add take_screenshot method

The `WebDriver:TakeScreenshot` command apparently also accepts an
`element` argument, but I am not exactly sure what value it expects.
---
 src/aiomarionette/__init__.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/src/aiomarionette/__init__.py b/src/aiomarionette/__init__.py
index 62d0455..6c26c90 100644
--- a/src/aiomarionette/__init__.py
+++ b/src/aiomarionette/__init__.py
@@ -174,6 +174,26 @@ class WebDriverBase(_BaseRPC, metaclass=abc.ABCMeta):
             'WebDriver:SwitchToWindow', handle=handle, focus=focus
         )
 
+    async def take_screenshot(
+        self,
+        *,
+        hash: bool = False,  # pylint: disable=redefined-builtin
+        full: bool = True,
+        scroll: bool = True,
+    ) -> str:
+        '''Take a screenshot of the current frame
+
+        :returns: Lossless PNG image encoded as a base-64 Unicode string
+        '''
+
+        res: Dict[str, str] = await self._send_message(
+            'WebDriver:TakeScreenshot',
+            full=full,
+            hash=hash,
+            scroll=scroll,
+        )
+        return res['value']
+
 
 class Marionette(WebDriverBase):
     '''Firefox Marionette session