From 14af2e7d1049f20fedac073d8e8515af58b8173a Mon Sep 17 00:00:00 2001
From: Thomas Walpole <twalpole@gmail.com>
Date: Fri, 27 Jan 2017 13:48:04 -0800
Subject: [PATCH] Support returning elements from evaluate_script

---
 lib/capybara/webkit/driver.rb | 20 ++++++++++++++++++--
 spec/driver_spec.rb           | 15 +++++++++++++++
 src/Evaluate.cpp              | 10 ++++++----
 src/Execute.cpp               |  5 +++--
 src/capybara.js               | 35 ++++++++++++++++++++++++++++-------
 5 files changed, 70 insertions(+), 15 deletions(-)

diff --git a/lib/capybara/webkit/driver.rb b/lib/capybara/webkit/driver.rb
index ce9dfa6..0609bbb 100644
--- a/lib/capybara/webkit/driver.rb
+++ b/lib/capybara/webkit/driver.rb
@@ -87,7 +87,8 @@ module Capybara::Webkit
     end
 
     def evaluate_script(script, *args)
-      @browser.evaluate_script(script, *encode_args(args))
+      result = @browser.evaluate_script(script, *encode_args(args))
+      decode_result(result)
     end
 
     def console_messages
@@ -419,11 +420,26 @@ module Capybara::Webkit
     def encode_args(args)
       args.map do |arg|
         if arg.is_a?(Capybara::Webkit::Node)
-          { ELEMENT: arg.native }.to_json
+          { 'element-581e-422e-8be1-884c4e116226' => arg.native }.to_json
         else
           arg.to_json
         end
       end
     end
+
+    def decode_result(result)
+      case result
+      when Array
+        result.map { |r| decode_result(r) }
+      when Hash
+        if element_ref = result['element-581e-422e-8be1-884c4e116226']
+          Capybara::Webkit::Node.new(self, element_ref, @browser)
+        else
+          result.each { |k,v| result[k] = decode_result(v) }
+        end
+      else
+        result
+      end
+    end
   end
 end
diff --git a/spec/driver_spec.rb b/spec/driver_spec.rb
index 9b09ec9..2721e3d 100644
--- a/spec/driver_spec.rb
+++ b/spec/driver_spec.rb
@@ -522,6 +522,21 @@ describe Capybara::Webkit::Driver do
       expect(result).to eq 1.5
     end
 
+    it "evaluates Javascript and returns an element" do
+      result = driver.evaluate_script(%<document.getElementById('greeting')>)
+      expect(result).to eq driver.find_css('#greeting').first
+    end
+
+    it "evaluates Javascript and returns a structure containing elements" do
+      result = driver.evaluate_script(%<({ 'a': document.getElementById('greeting'), 'b': { 'c': document.querySelectorAll('#greeting, #checktest') } })>)
+      expect(result).to eq({
+        'a' => driver.find_css('#greeting').first,
+        'b' => {
+          'c' => driver.find_css('#greeting, #checktest')
+        }
+      })
+    end
+
     it "evaluates Javascript and returns null" do
       result = driver.evaluate_script(%<(function () {})()>)
       expect(result).to eq nil
diff --git a/src/Evaluate.cpp b/src/Evaluate.cpp
index 7ab80e2..61bcd94 100644
--- a/src/Evaluate.cpp
+++ b/src/Evaluate.cpp
@@ -18,12 +18,14 @@ void Evaluate::start() {
   QString eval_script = QString("(function(){"
                            "   for(var i=0; i<arguments.length; i++) {"
                            "     arguments[i] = JSON.parse(arguments[i]);"
-                           "     if (arguments[i]['ELEMENT']) {"
-                           "       arguments[i] = Capybara.getNode(arguments[i]['ELEMENT']);"
+                           "     var elem_id;"
+                           "     if (elem_id = arguments[i]['element-581e-422e-8be1-884c4e116226']) {"
+                           "       arguments[i] = Capybara.getNode(elem_id);"
                            "     };"
                            "   };"
-                           "   return eval(\"%1\");"
-                           "   }).apply(null, %2)").arg(script.replace("\"","\\\"").remove("\n"), jsonArgs);
+                           "   var _result = eval(\"%1\");"
+                           "   return Capybara.wrapResult(_result);"
+                           " }).apply(null, %2)").arg(script.replace("\"","\\\"").remove("\n"), jsonArgs);
   QObject invocation_stub;
   invocation_stub.setProperty("allowUnattached", false);
   page()->currentFrame()->addToJavaScriptWindowObject("CapybaraInvocation", &invocation_stub);
diff --git a/src/Execute.cpp b/src/Execute.cpp
index 325007a..a4c54f6 100644
--- a/src/Execute.cpp
+++ b/src/Execute.cpp
@@ -16,8 +16,9 @@ void Execute::start() {
   QString script = QString("(function(){"
                            "   for(var i=0; i<arguments.length; i++) {"
                            "     arguments[i] = JSON.parse(arguments[i]);"
-                           "     if (arguments[i]['ELEMENT']) {"
-                           "       arguments[i] = Capybara.getNode(arguments[i]['ELEMENT']);"
+                           "     var elem_id;"
+                           "     if (elem_id = arguments[i]['element-581e-422e-8be1-884c4e116226']) {"
+                           "       arguments[i] = Capybara.getNode(elem_id);"
                            "     };"
                            "   };"
                            "   %1 }).apply(null, %2); 'success'").arg(arguments()[0], jsonArgs);
diff --git a/src/capybara.js b/src/capybara.js
index 623a3bf..5302196 100644
--- a/src/capybara.js
+++ b/src/capybara.js
@@ -37,9 +37,7 @@ Capybara = {
     var node;
     var results = [];
     while (node = iterator.iterateNext()) {
-      this.nextIndex++;
-      this.nodes[this.nextIndex] = node;
-      results.push(this.nextIndex);
+      results.push(this.registerNode(node));
     }
     return results.join(",");
   },
@@ -48,9 +46,7 @@ Capybara = {
     var elements = reference.querySelectorAll(selector);
     var results = [];
     for (var i = 0; i < elements.length; i++) {
-      this.nextIndex++;
-      this.nodes[this.nextIndex] = elements[i];
-      results.push(this.nextIndex);
+      results.push(this.registerNode(elements[i]));
     }
     return results.join(",");
   },
@@ -477,9 +473,34 @@ Capybara = {
     mouseTrigger('mousemove', options);
     mouseTrigger('mouseup', options);
   },
-
+  registerNode: function(node) {
+    this.nextIndex++;
+    this.nodes[this.nextIndex] = node;
+    return this.nextIndex;
+  },
   equals: function(index, targetIndex) {
     return this.getNode(index) === this.getNode(targetIndex);
+  },
+  _visitedObjects: [],
+  wrapResult: function(arg) {
+    if (this._visitedObjects.indexOf(arg) >= 0) { return '(cyclic structure)'; }
+    if (arg instanceof NodeList) { arg = Array.prototype.slice.call(arg, 0); }
+    if (Array.isArray(arg)) {
+      for(var _j = 0; _j < arg.length; _j++) {
+        arg[_j] = this.wrapResult(arg[_j]);
+      }
+    } else if (arg && arg.nodeType == 1 && arg.tagName) {
+      return {'element-581e-422e-8be1-884c4e116226': this.registerNode(arg) }
+    } else if (arg === null) {
+      return undefined;
+    } else if ( typeof arg == 'object' ) {
+      this._visitedObjects.push(arg);
+      for(var _k in arg){
+        arg[_k] = this.wrapResult(arg[_k]);
+      }
+      this._visitedObjects.pop();
+    }
+    return arg;
   }
 };