Wrap now returns a proper wrapped object instead just overloading the…

… load-function

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
 

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) {

test.js

@@ -93,30 +93,47 @@ 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()
     })
   })
 
   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']
 
     $('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())
+    })
   })
 })