Implement getIteratorFn function to get iterator function for iterabl…

…e Object.

Update isIterable function to leverage getIteratorFn.
Rewrite flatten function to leverage getIteratorFn.
Add unit test for flatten function.
Add unit tests for arbitrary iterable children.
Support for Symbol.iterator and @@iterator iterables.

packages/enzyme-adapter-utils/src/Utils.js

@@ -92,24 +92,55 @@ export function nodeTypeFromType(type) {
   return 'function';
 }
 
-function isIterable(obj) {
-  return (
-    obj != null &&
-    typeof Symbol === 'function' &&
-    typeof Symbol.iterator === 'symbol' &&
-    typeof obj[Symbol.iterator] === 'function'
+function getIteratorFn(obj) {
+  const iteratorFn = obj && (
+    (
+      typeof Symbol === 'function' &&
+      typeof Symbol.iterator === 'symbol' &&
+      obj[Symbol.iterator]
+    ) ||
+    obj['@@iterator']
   );
+
+  if (typeof iteratorFn === 'function') {
+    return iteratorFn;
+  }
+
+  return undefined;
+}
+
+function isIterable(obj) {
+  return Boolean(getIteratorFn(obj));
 }
 
 export function isArrayLike(obj) {
   return Array.isArray(obj) || (isIterable(obj) && typeof obj !== 'string');
 }
 
 export function flatten(arrs) {
-  return arrs.reduce(
-    (flattened, item) => flattened.concat(isArrayLike(item) ? flatten([...item]) : item),
-    [],
-  );
+  let flatArrs = [];
+
+  const iteratorFn = getIteratorFn(arrs);
+  const iterator = iteratorFn.call(arrs);
+
+  let step = iterator.next();
+
+  while (!step.done) {
+    const item = step.value;
+    let flatItem;
+
+    if (isArrayLike(item)) {
+      flatItem = flatten(item);
+    } else {
+      flatItem = item;
+    }
+
+    flatArrs = flatArrs.concat(flatItem);
+
+    step = iterator.next();
+  }
+
+  return flatArrs;
 }
 
 export function elementToTree(el) {
@@ -125,7 +156,7 @@ export function elementToTree(el) {
   const { children } = props;
   let rendered = null;
   if (isArrayLike(children)) {
-    rendered = flatten([...children], true).map(elementToTree);
+    rendered = flatten(children).map(elementToTree);
   } else if (typeof children !== 'undefined') {
     rendered = elementToTree(children);
   }

packages/enzyme-test-suite/test/RSTTraversal-spec.jsx

@@ -225,6 +225,115 @@ describe('RSTTraversal', () => {
       expect(nodes).to.deep.equal([node, divA, divB, divC, divD]);
     });
 
+    describe('support for arbitrary iterable children', () => {
+      const makeDivIterator = (lowerBound, upperBound) => {
+        let counter = lowerBound;
+
+        return {
+          next() {
+            if (counter < upperBound) {
+              const key = String.fromCharCode('a'.charCodeAt(0) + counter);
+
+              const nextValue = {
+                value: <div key={key} />,
+                done: false,
+              };
+
+              counter += 1;
+
+              return nextValue;
+            }
+
+            return { done: true };
+          },
+        };
+      };
+
+      it('should handle iterable with Symbol.iterator property children', () => {
+        const spy = sinon.spy();
+
+        const iterableChildren = { [Symbol.iterator]: () => makeDivIterator(0, 2) };
+
+        const divA = $(<div key="a" />);
+        const divB = $(<div key="b" />);
+        const node = $((
+          <div>
+            {iterableChildren}
+          </div>
+        ));
+
+        treeForEach(node, spy);
+        expect(spy.callCount).to.equal(3);
+        const nodes = spy.args.map(arg => arg[0]);
+        expect(nodes).to.deep.equal([node, divA, divB]);
+      });
+
+      it('should handle iterable with Symbol.iterator property siblings', () => {
+        const spy = sinon.spy();
+
+        const iterableChildren1 = { [Symbol.iterator]: () => makeDivIterator(0, 2) };
+        const iterableChildren2 = { [Symbol.iterator]: () => makeDivIterator(2, 4) };
+
+        const divA = $(<div key="a" />);
+        const divB = $(<div key="b" />);
+        const divC = $(<div key="c" />);
+        const divD = $(<div key="d" />);
+        const node = $((
+          <div>
+            {iterableChildren1}
+            {iterableChildren2}
+          </div>
+        ));
+
+        treeForEach(node, spy);
+        expect(spy.callCount).to.equal(5);
+        const nodes = spy.args.map(arg => arg[0]);
+        expect(nodes).to.deep.equal([node, divA, divB, divC, divD]);
+      });
+
+      it('should handle iterable with @@iterator property children', () => {
+        const spy = sinon.spy();
+
+        const legacyIterableChildren = { '@@iterator': () => makeDivIterator(0, 2) };
+
+        const divA = $(<div key="a" />);
+        const divB = $(<div key="b" />);
+        const node = $((
+          <div>
+            {legacyIterableChildren}
+          </div>
+        ));
+
+        treeForEach(node, spy);
+        expect(spy.callCount).to.equal(3);
+        const nodes = spy.args.map(arg => arg[0]);
+        expect(nodes).to.deep.equal([node, divA, divB]);
+      });
+
+      it('should handle iterable with @@iterator property siblings', () => {
+        const spy = sinon.spy();
+
+        const legacyIterableChildren1 = { '@@iterator': () => makeDivIterator(0, 2) };
+        const legacyIterableChildren2 = { '@@iterator': () => makeDivIterator(2, 4) };
+
+        const divA = $(<div key="a" />);
+        const divB = $(<div key="b" />);
+        const divC = $(<div key="c" />);
+        const divD = $(<div key="d" />);
+        const node = $((
+          <div>
+            {legacyIterableChildren1}
+            {legacyIterableChildren2}
+          </div>
+        ));
+
+        treeForEach(node, spy);
+        expect(spy.callCount).to.equal(5);
+        const nodes = spy.args.map(arg => arg[0]);
+        expect(nodes).to.deep.equal([node, divA, divB, divC, divD]);
+      });
+    });
+
     it('should not get trapped from empty strings', () => {
       const spy = sinon.spy();
       const node = $((

packages/enzyme-test-suite/test/Utils-spec.jsx

@@ -7,6 +7,7 @@ import {
   displayNameOfNode,
 } from 'enzyme/build/Utils';
 import {
+  flatten,
   mapNativeEventNames,
   propFromEvent,
 } from 'enzyme-adapter-utils';
@@ -542,4 +543,12 @@ describe('Utils', () => {
       expectEqualArrays(childrenToSimplifiedArray(children), simplified);
     });
   });
+
+  describe('flatten', () => {
+    it('should recursively flatten a nested iterable structure', () => {
+      const nested = [1, [2, [3, [4]], 5], 6, [7, [8, 9]], 10];
+      const flat = flatten(nested);
+      expect(flat).to.deep.equal([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+    });
+  });
 });