Enable prop/state capture in JSX decorators (#917)

examples/jsx-fixtures/components/WelcomeMessage/cosmos.decorator.js

@@ -0,0 +1,20 @@
+// @flow
+
+import React from 'react';
+import { FixtureCapture } from 'react-cosmos-fixture';
+
+export default ({ children }: { children: React$Node }) => (
+  <FixtureCapture decoratorId="bgDecorator">
+    <BgDecorator backgroundColor="lightgrey">{children}</BgDecorator>
+  </FixtureCapture>
+);
+
+function BgDecorator({
+  children,
+  backgroundColor
+}: {
+  children: React$Node,
+  backgroundColor: string
+}) {
+  return <div style={{ backgroundColor }}>{children}</div>;
+}

packages/react-cosmos-fixture/package.json

@@ -8,6 +8,7 @@
   "dependencies": {
     "@babel/runtime-corejs2": "^7.2.0",
     "lodash": "^4.17.11",
+    "memoize-one": "^5.0.0",
     "react-cosmos-shared2": "^4.7.0-22",
     "react-is": "^16.6.3",
     "socket.io-client": "^2.2.0"

packages/react-cosmos-fixture/src/FixtureCapture/attachChildRefs/index.js

@@ -1,15 +1,14 @@
 // @flow
 
 import { cloneElement, Component } from 'react';
-import { setElementAtPath } from '../childrenTree';
+import { setElementAtPath } from '../nodeTree';
 import { findRelevantElementPaths } from '../findRelevantElementPaths';
 import { compose } from './compose';
 import { isRefSupported } from './isRefSupported';
 import { createRefHandler } from './createRefHandler';
 
-import type { ElementRef, Ref } from 'react';
+import type { Node, ElementRef, Ref } from 'react';
 import type { FixtureDecoratorId } from 'react-cosmos-shared2/fixtureState';
-import type { Children } from '../childrenTree';
 import type { ComponentRef } from '../shared';
 
 type RefWrapper = {
@@ -31,20 +30,20 @@ const refHandlers: WeakMap<
 > = new WeakMap();
 
 export function attachChildRefs({
-  children,
+  node,
   onRef,
   decoratorElRef,
   decoratorId
 }: {
-  children: Children,
+  node: Node,
   onRef: (elPath: string, elRef: ?ComponentRef) => mixed,
   decoratorElRef: ElementRef<typeof Component>,
   decoratorId: FixtureDecoratorId
 }) {
-  const elPaths = findRelevantElementPaths(children);
+  const elPaths = findRelevantElementPaths(node);
 
-  return elPaths.reduce((extendedChildren, elPath): Children => {
-    return setElementAtPath(extendedChildren, elPath, element => {
+  return elPaths.reduce((extendedNode, elPath): Node => {
+    return setElementAtPath(extendedNode, elPath, element => {
       if (!isRefSupported(element.type)) {
         return element;
       }
@@ -59,7 +58,7 @@ export function attachChildRefs({
         })
       });
     });
-  }, children);
+  }, node);
 }
 
 export function deleteRefHandler(

packages/react-cosmos-fixture/src/FixtureCapture/extendChildPropsWithFixtureState.js → packages/react-cosmos-fixture/src/FixtureCapture/extendPropsWithFixtureState.js

@@ -5,23 +5,23 @@ import {
   extendObjWithValues,
   findCompFixtureState
 } from 'react-cosmos-shared2/fixtureState';
-import { setElementAtPath } from './childrenTree';
+import { setElementAtPath } from './nodeTree';
 import { findRelevantElementPaths } from './findRelevantElementPaths';
 
+import type { Node } from 'react';
 import type {
   FixtureDecoratorId,
   FixtureState
 } from 'react-cosmos-shared2/fixtureState';
-import type { Children } from './childrenTree';
 
-export function extendChildPropsWithFixtureState(
-  children: Children,
+export function extendPropsWithFixtureState(
+  node: Node,
   fixtureState: null | FixtureState,
   decoratorId: FixtureDecoratorId
-): Children {
-  const elPaths = findRelevantElementPaths(children);
+): Node {
+  const elPaths = findRelevantElementPaths(node);
 
-  return elPaths.reduce((extendedChildren, elPath): Children => {
+  return elPaths.reduce((extendedChildren, elPath): Node => {
     const compFxState = findCompFixtureState(fixtureState, decoratorId, elPath);
 
     return setElementAtPath(extendedChildren, elPath, element => {
@@ -54,7 +54,7 @@ export function extendChildPropsWithFixtureState(
         key: getElRenderKey(elPath, compFxState.renderKey)
       };
     });
-  }, children);
+  }, node);
 }
 
 function getElRenderKey(elPath, renderKey) {

packages/react-cosmos-fixture/src/FixtureCapture/findRelevantElementPaths.js

@@ -1,14 +1,14 @@
 // @flow
 
-import { findElementPaths, getExpectedElementAtPath } from './childrenTree';
+import { findElementPaths, getExpectedElementAtPath } from './nodeTree';
 
-import type { Children } from './childrenTree';
+import type { Node } from 'react';
 
-export function findRelevantElementPaths(children: Children): string[] {
-  const elPaths = findElementPaths(children);
+export function findRelevantElementPaths(node: Node): string[] {
+  const elPaths = findElementPaths(node);
 
   return elPaths.filter(elPath => {
-    const { type } = getExpectedElementAtPath(children, elPath);
+    const { type } = getExpectedElementAtPath(node, elPath);
 
     // TODO: Make this customizable
     if (type === 'string') {

packages/react-cosmos-fixture/src/FixtureCapture/index.js

@@ -15,12 +15,12 @@ import { FixtureContext } from '../FixtureContext';
 import {
   getElementAtPath,
   getExpectedElementAtPath,
-  areChildrenEqual
-} from './childrenTree';
+  areNodesEqual
+} from './nodeTree';
 import { getComponentName } from './getComponentName';
 import { getElementRefType } from './getElementRefType';
 import { findRelevantElementPaths } from './findRelevantElementPaths';
-import { extendChildPropsWithFixtureState } from './extendChildPropsWithFixtureState';
+import { extendPropsWithFixtureState } from './extendPropsWithFixtureState';
 import {
   attachChildRefs,
   deleteRefHandler,
@@ -83,11 +83,7 @@ class FixtureCaptureInner extends Component<InnerProps> {
     const { children, decoratorId, fixtureState } = this.props;
 
     return attachChildRefs({
-      children: extendChildPropsWithFixtureState(
-        children,
-        fixtureState,
-        decoratorId
-      ),
+      node: extendPropsWithFixtureState(children, fixtureState, decoratorId),
       onRef: this.handleRef,
       decoratorElRef: this,
       decoratorId
@@ -145,7 +141,7 @@ class FixtureCaptureInner extends Component<InnerProps> {
     const { children, decoratorId, fixtureState } = this.props;
 
     // Children change when the fixture is updated at runtime (eg. via HMR)
-    if (!areChildrenEqual(nextProps.children, children)) {
+    if (!areNodesEqual(nextProps.children, children)) {
       return true;
     }
 
@@ -200,7 +196,7 @@ class FixtureCaptureInner extends Component<InnerProps> {
       // the fixture.
       !compFxState.props ||
       // b) mocked props from fixture elemented changed (likely via HMR).
-      !areChildrenEqual(childEl, getElementAtPath(prevProps.children, elPath))
+      !areNodesEqual(childEl, getElementAtPath(prevProps.children, elPath))
     ) {
       this.updateFixtureState({ elPath, props: childEl.props });
     }

packages/react-cosmos-fixture/src/FixtureCapture/childrenTree/__tests__/find.js → packages/react-cosmos-fixture/src/FixtureCapture/nodeTree/__tests__/find.js

@@ -75,10 +75,6 @@ it('finds nested paths', () => {
   });
 });
 
-it('finds no paths on function children', () => {
-  expect(findElementPaths(() => null)).toEqual([]);
-});
-
 it('only finds paths outside function children', () => {
   const Comp = () => null;
   expect(

packages/react-cosmos-fixture/src/FixtureCapture/nodeTree/areNodesEqual.js

@@ -0,0 +1,33 @@
+// @flow
+
+import { isEqual, pick, mapValues } from 'lodash';
+import { isElement } from 'react-is';
+
+import type { Node, Element } from 'react';
+
+export function areNodesEqual(a: Node, b: Node): boolean {
+  return isEqual(stripInternalElementAttrs(a), stripInternalElementAttrs(b));
+}
+
+// Don't compare private element attrs like _owner and _store, which hold
+// internal details and have auto increment-type attrs
+function stripInternalElementAttrs(node: mixed) {
+  if (Array.isArray(node)) {
+    return node.map(n => stripInternalElementAttrs(n));
+  }
+
+  if (isElement(node)) {
+    // $FlowFixMe Flow can't get cues from react-is package
+    const el: Element<any> = node;
+
+    return {
+      ...pick(el, 'type', 'key', 'ref'),
+      // children and other props can contain Elements
+      props: mapValues(el.props, propValue =>
+        stripInternalElementAttrs(propValue)
+      )
+    };
+  }
+
+  return node;
+}

packages/react-cosmos-fixture/src/FixtureCapture/nodeTree/findElementPaths.js

@@ -0,0 +1,39 @@
+// @flow
+
+import { flatten } from 'lodash';
+import { isElement } from 'react-is';
+import { Fragment } from 'react';
+import { isRootPath } from './shared';
+
+import type { Node, Element } from 'react';
+
+export function findElementPaths(node: Node, curPath: string = ''): string[] {
+  if (Array.isArray(node)) {
+    return flatten(
+      node.map((child, idx) => findElementPaths(child, `${curPath}[${idx}]`))
+    );
+  }
+
+  if (!isElement(node)) {
+    // At this point the node can be null, boolean, string, number, Portal, etc.
+    // https://github.com/facebook/flow/blob/172d28f542f49bbc1e765131c9dfb9e31780f3a2/lib/react.js#L13-L20
+    return [];
+  }
+
+  // $FlowFixMe Flow can't get cues from react-is package
+  const element: Element<any> = node;
+  const { children } = element.props;
+
+  const childElPaths =
+    // Props of elements returned by render functions can't be read here
+    typeof children !== 'function'
+      ? findElementPaths(children, getChildrenPath(curPath))
+      : [];
+
+  // Ignore Fragment elements, but include their children
+  return element.type === Fragment ? childElPaths : [curPath, ...childElPaths];
+}
+
+function getChildrenPath(curPath) {
+  return isRootPath(curPath) ? 'props.children' : `${curPath}.props.children`;
+}