diff --git a/src/ReactWrapper.jsx b/src/ReactWrapper.jsx index f663a7d90..c3a093ba9 100644 --- a/src/ReactWrapper.jsx +++ b/src/ReactWrapper.jsx @@ -63,7 +63,7 @@ function filterWhereUnwrapped(wrapper, predicate) { /** * @class ReactWrapper */ -export default class ReactWrapper { +class ReactWrapper { constructor(nodes, root, options = {}) { if (!global.window && !global.document) { @@ -107,26 +107,6 @@ export default class ReactWrapper { ); } - /** - * Makes a wrapper iterable, which is useful when using destructive - * assignment with `find()` - * @example - * const wrapper = shallow() - * const [fistLink, secondLink] = wrapper.find('a') - * @returns {Object} - */ - [ITERATOR_SYMBOL]() { - let index = 0; - return { - next: () => { - if (index >= this.nodes.length) { - return { done: true }; - } - return { done: false, value: this.nodes[index++] }; - }, - }; - } - /** * If the root component contained a ref, you can access it here * and get a wrapper around it. @@ -929,3 +909,12 @@ export default class ReactWrapper { unmountComponentAtNode(this.options.attachTo); } } + + +if (ITERATOR_SYMBOL) { + ReactWrapper.prototype[ITERATOR_SYMBOL] = function iterator() { + return this.nodes[ITERATOR_SYMBOL](); + }; +} + +export default ReactWrapper; diff --git a/src/ShallowWrapper.js b/src/ShallowWrapper.js index 6d419c7b3..6fcead563 100644 --- a/src/ShallowWrapper.js +++ b/src/ShallowWrapper.js @@ -64,7 +64,7 @@ function filterWhereUnwrapped(wrapper, predicate) { /** * @class ShallowWrapper */ -export default class ShallowWrapper { +class ShallowWrapper { constructor(nodes, root, options = {}) { if (!root) { @@ -104,26 +104,6 @@ export default class ShallowWrapper { this.complexSelector = new ComplexSelector(buildPredicate, findWhereUnwrapped, childrenOfNode); } - /** - * Makes a wrapper iterable, which is useful when using destructive - * assignment with `find()` - * @example - * const wrapper = shallow() - * const [fistLink, secondLink] = wrapper.find('a') - * @returns {Object} - */ - [ITERATOR_SYMBOL]() { - let index = 0; - return { - next: () => { - if (index >= this.nodes.length) { - return { done: true }; - } - return { done: false, value: this.nodes[index++] }; - }, - }; - } - /** * Gets the instance of the component being rendered as the root node passed into `shallow()`. * @@ -1003,3 +983,11 @@ export default class ShallowWrapper { }); } } + +if (ITERATOR_SYMBOL) { + ShallowWrapper.prototype[ITERATOR_SYMBOL] = function iterator() { + return this.nodes[ITERATOR_SYMBOL](); + }; +} + +export default ShallowWrapper; diff --git a/src/Utils.js b/src/Utils.js index dc956f47e..5187321c3 100644 --- a/src/Utils.js +++ b/src/Utils.js @@ -12,9 +12,7 @@ import { REACT15, } from './version'; -export const ITERATOR_SYMBOL = ( - typeof Symbol === 'function' && Symbol.iterator -) || '@@iterator'; +export const ITERATOR_SYMBOL = typeof Symbol === 'function' && Symbol.iterator; function internalInstanceKey(node) { return Object.keys(Object(node)).filter(key => key.match(/^__reactInternalInstance\$/))[0];