From af85914c15811bf6aa1dfe394b3a8ea158a2a71a Mon Sep 17 00:00:00 2001
From: Thomas Watson Steen <w@tson.dk>
Date: Tue, 30 Jun 2015 12:35:36 -0400
Subject: [PATCH] Wrap now returns a proper wrapped object instead just
 overloading the load-function

---
 README.md |  7 ++++---
 index.js  | 26 ++++++++++++++++++++------
 test.js   | 27 ++++++++++++++++++++++-----
 3 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/README.md b/README.md
index fc33830..56708da 100644
--- a/README.md
+++ b/README.md
@@ -38,9 +38,10 @@ var $ = cheerio.load('<div>foo</div><div>bar</div>')
 $('div:first').text() // => 'foo'
 ```
 
-Note that this will only work if the HTML is loaded using the `.load()`
-function as seen above (see [issue
-1](https://github.com/watson/cheerio-advanced-selectors/issues/1)).
+_**Gotcha:** The result returned from `.load()` isn't a cheerio object
+but a custom function used to wrap the cheerio-advanced-selector logic
+(see [issue
+2](https://github.com/watson/cheerio-advanced-selectors/issues/2))._
 
 ## Advanced usage
 
diff --git a/index.js b/index.js
index d404f54..5cd9113 100644
--- a/index.js
+++ b/index.js
@@ -1,19 +1,33 @@
 'use strict'
 
+var util = require('util')
+
 var splitter = /^(.*?)(?:\:(eq|first|last)(?:\((\d+)\))?)(.*)/
 
-exports.wrap = function (cheerio) {
-  var load = cheerio.load
+exports.wrap = function (Cheerio) {
+  var CheerioAdv = function (selector, context, root, opts) {
+    if (!(this instanceof CheerioAdv)) return new CheerioAdv(selector, context, root, opts)
+
+    if (typeof selector === 'string' && splitter.test(selector)) {
+      var steps = split(selector)
+      var cursor = Cheerio(steps.shift(), context, root, opts)
+      return execSteps(cursor, steps)
+    }
+
+    return Cheerio.apply(Cheerio, arguments)
+  }
+
+  util.inherits(CheerioAdv, Cheerio)
 
-  cheerio.load = function () {
-    var $ = load.apply(cheerio, arguments)
+  CheerioAdv.load = function () {
+    var $ = Cheerio.load.apply(Cheerio, arguments)
     return function (selector, context, root) {
       if (typeof selector === 'string') return exports.find($, selector, context, root)
-      return $.apply(cheerio, arguments)
+      return $.apply(Cheerio, arguments)
     }
   }
 
-  return cheerio
+  return CheerioAdv
 }
 
 exports.find = function ($, selector, context, root) {
diff --git a/test.js b/test.js
index 31ae59b..ee1a649 100644
--- a/test.js
+++ b/test.js
@@ -93,14 +93,12 @@ test('#find()', function (t) {
   })
 })
 
-test('#wrap()', function (t) {
+test('#wrap() -> #load()', function (t) {
   testCases.forEach(function (testCase) {
     t.test(testCase[0], function (t) {
-      var load = cheerio.load
       var wrapped = cheerioAdv.wrap(cheerio)
       var $ = wrapped.load(testCase[1])
       t.equal($(testCase[2]).text(), testCase[3])
-      cheerio.load = load
       t.end()
     })
   })
@@ -108,7 +106,6 @@ test('#wrap()', function (t) {
   t.test('this as selector', function (t) {
     t.plan(2)
 
-    var load = cheerio.load
     var wrapped = cheerioAdv.wrap(cheerio)
     var $ = wrapped.load('<div>foo</div><div>bar</div>')
     var results = ['foo', 'bar']
@@ -116,7 +113,27 @@ test('#wrap()', function (t) {
     $('div').each(function () {
       t.equal($(this).text(), results.shift())
     })
+  })
+})
+
+test('#wrap() -> #()', function (t) {
+  testCases.forEach(function (testCase) {
+    t.test(testCase[0], function (t) {
+      var wrapped = cheerioAdv.wrap(cheerio)
+      t.equal(wrapped(testCase[2], testCase[1]).text(), testCase[3])
+      t.end()
+    })
+  })
 
-    cheerio.load = load
+  t.test('this as selector', function (t) {
+    t.plan(2)
+
+    var wrapped = cheerioAdv.wrap(cheerio)
+    var html = '<div>foo</div><div>bar</div>'
+    var results = ['foo', 'bar']
+
+    wrapped('div', html).each(function () {
+      t.equal(wrapped(this).text(), results.shift())
+    })
   })
 })