Merge pull request #1334 from Inucoder/better_support_for_iterable_ch…

…ildren

`enzyme-adapter-utils`: [New] Better support for iterable children

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

@@ -92,24 +92,64 @@ 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 !!getIteratorFn(obj);
 }
 
 export function isArrayLike(obj) {
-  return Array.isArray(obj) || (isIterable(obj) && typeof obj !== 'string');
+  return Array.isArray(obj) || (typeof obj !== 'string' && isIterable(obj));
 }
 
 export function flatten(arrs) {
-  return arrs.reduce(
-    (flattened, item) => flattened.concat(isArrayLike(item) ? flatten([...item]) : item),
-    [],
-  );
+  // optimize for the most common case
+  if (Array.isArray(arrs)) {
+    return arrs.reduce(
+      (flatArrs, item) => flatArrs.concat(isArrayLike(item) ? flatten(item) : item),
+      [],
+    );
+  }
+
+  // fallback for arbitrary iterable children
+  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 +165,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,116 @@ describe('RSTTraversal', () => {
       expect(nodes).to.deep.equal([node, divA, divB, divC, divD]);
     });
 
+    describe('support for arbitrary iterable children', () => {
+      const makeDivIterator = (lowerBound, upperBound) => {
+        const baseCode = 'a'.charCodeAt(0);
+        let counter = lowerBound;
+
+        return {
+          next() {
+            if (counter < upperBound) {
+              const key = String.fromCharCode(baseCode + 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]);
+    });
+  });
 });