diff --git a/index.js b/index.js index a847c8f..4943ce2 100644 --- a/index.js +++ b/index.js @@ -46,7 +46,7 @@ function addWith(obj, src, exclude) { src = '(function (' + vars.join(', ') + ') {' + src + - '}(' + inputVars.join(',') + '))' + '}.call(this' + inputVars.map(function (v) { return ',' + v; }).join('') + '))' return ';' + declareLocal + ';' + unwrapReturns(src, result) + ';' } @@ -77,12 +77,17 @@ function unwrapReturns(src, result) { var originalSource = src var hasReturn = false var ast = uglify.parse(src) + var ref src = src.split('') - if (ast.body.length !== 1 || ast.body[0].TYPE !== 'SimpleStatement' || - ast.body[0].body.TYPE !== 'Call' || ast.body[0].body.expression.TYPE !== 'Function') + // get a reference to the function that was inserted to add an inner context + if ((ref = ast.body).length !== 1 + || (ref = ref[0]).TYPE !== 'SimpleStatement' + || (ref = ref.body).TYPE !== 'Call' + || (ref = ref.expression).TYPE !== 'Dot' || ref.property !== 'call' + || (ref = ref.expression).TYPE !== 'Function') throw new Error('AST does not seem to represent a self-calling function') - var fn = ast.body[0].body.expression + var fn = ref var walker = new uglify.TreeWalker(visitor) function visitor(node, descend) { diff --git a/test/index.js b/test/index.js index a0bbb45..b662056 100644 --- a/test/index.js +++ b/test/index.js @@ -126,4 +126,18 @@ describe('addWith("obj || {}", "return foo")', function () { assert(Function('obj', src + ';return "ding"')({}) === 'ding') done() }) -}) \ No newline at end of file +}) + +describe('addWith("obj || {}", "return this[foo]")', function () { + it('keeps reference to this', function (done) { + var src = addWith('obj || {}', 'return this[foo]') + // obj.bar = obj.foo + var obj = { + foo: 'bar', + bar: 'ding', + fn: Function('obj', src) + } + assert(obj.fn(obj) === 'ding') + done() + }) +})