Convert element to object from execute result

lib/WebDriver/Session.php

@@ -34,8 +34,6 @@
  * @method void forward() Navigates forward in the browser history, if possible.
  * @method void back() Navigates backward in the browser history, if possible.
  * @method void refresh() Refresh the current page.
- * @method mixed execute($jsonScript) Inject a snippet of JavaScript into the page for execution in the context of the currently selected frame. (synchronous)
- * @method mixed execute_async($jsonScript) Inject a snippet of JavaScript into the page for execution in the context of the currently selected frame. (asynchronous)
  * @method string screenshot() Take a screenshot of the current page.
  * @method array getCookie() Retrieve all cookies visible to the current page.
  * @method array postCookie($jsonCookie) Set a cookie.
@@ -433,4 +431,77 @@ protected function getElementPath($elementId)
     {
         return sprintf('%s/element/%s', $this->url, $elementId);
     }
+
+    /**
+     * @param array $args
+     *
+     * @return array
+     */
+    private function prepareScriptArguments(array $args)
+    {
+        foreach ($args as $k => $v) {
+            if ($v instanceof Element) {
+                $args[$k] = [Container::WEBDRIVER_ELEMENT_ID => $v->getID()];
+            } elseif (is_array($v)) {
+                $args[$k] = $this->prepareScriptArguments($v);
+            }
+        }
+
+        return $args;
+    }
+
+    /**
+     * @param mixed $data
+     *
+     * @return mixed
+     */
+    private function webDriverElementRecursive($data)
+    {
+        $element = $this->webDriverElement($data);
+        if ($element !== null) {
+            return $element;
+        } elseif (is_array($data)) {
+            foreach ($data as $k => $v) {
+                $data[$k] = $this->webDriverElementRecursive($v);
+            }
+        }
+
+        return $data;
+    }
+
+    /**
+     * Inject a snippet of JavaScript into the page for execution in the context of the currently selected frame. (synchronous)
+     *
+     * @param array{script: string, args: array} $jsonScript
+     *
+     * @return mixed
+     */
+    public function execute(array $jsonScript)
+    {
+        if (isset($jsonScript['args'])) {
+            $jsonScript['args'] = $this->prepareScriptArguments($jsonScript['args']);
+        }
+
+        $result = parent::execute($jsonScript);
+
+        return $this->webDriverElementRecursive($result);
+    }
+
+    /**
+     * Inject a snippet of JavaScript into the page for execution in the context of the currently selected frame. (asynchronous)
+     *
+     * @param array{script: string, args: array} $jsonScript
+     *
+     * @return mixed
+     */
+    public function execute_async(array $jsonScript)
+    {
+        if (isset($jsonScript['args'])) {
+            $jsonScript['args'] = $this->prepareScriptArguments($jsonScript['args']);
+        }
+
+        $result = parent::execute_async($jsonScript);
+
+        return $this->webDriverElementRecursive($result);
+    }
 }