Skip to content

Commit b283d75

Browse files
awearygaearon
aweary
authored and
gaearon
committed
Support React.memo in ReactShallowRenderer (#14816)
* Support React.memo in ReactShallowRenderer ReactShallowRenderer uses element.type frequently, but with React.memo elements the actual type is element.type.type. This updates ReactShallowRenderer so it uses the correct element type for Memo components and also validates the inner props for the wrapped components. * Allow Rect.memo to prevent re-renders * Support memo(forwardRef()) * Dont call memo comparison function on initial render * Fix test * Small tweaks
1 parent f0621fe commit b283d75

File tree

3 files changed

+1717
-39
lines changed

3 files changed

+1717
-39
lines changed

‎packages/react-test-renderer/src/ReactShallowRenderer.js‎

Lines changed: 86 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
*/
99

1010
import React from 'react';
11-
import {isForwardRef} from 'react-is';
11+
import {isForwardRef, isMemo, ForwardRef} from 'react-is';
1212
import describeComponentFrame from 'shared/describeComponentFrame';
1313
import getComponentName from 'shared/getComponentName';
1414
import shallowEqual from 'shared/shallowEqual';
@@ -500,7 +500,8 @@ class ReactShallowRenderer {
500500
element.type,
501501
);
502502
invariant(
503-
isForwardRef(element) || typeof element.type === 'function',
503+
isForwardRef(element) ||
504+
(typeof element.type === 'function' || isMemo(element.type)),
504505
'ReactShallowRenderer render(): Shallow rendering works only with custom ' +
505506
'components, but the provided element type was `%s`.',
506507
Array.isArray(element.type)
@@ -514,22 +515,36 @@ class ReactShallowRenderer {
514515
return;
515516
}
516517

518+
const elementType = isMemo(element.type) ? element.type.type : element.type;
519+
const previousElement = this._element;
520+
517521
this._rendering = true;
518522
this._element = element;
519-
this._context = getMaskedContext(element.type.contextTypes, context);
523+
this._context = getMaskedContext(elementType.contextTypes, context);
524+
525+
// Inner memo component props aren't currently validated in createElement.
526+
if (isMemo(element.type) && elementType.propTypes) {
527+
currentlyValidatingElement = element;
528+
checkPropTypes(
529+
elementType.propTypes,
530+
element.props,
531+
'prop',
532+
getComponentName(elementType),
533+
getStackAddendum,
534+
);
535+
}
520536

521537
if (this._instance) {
522-
this._updateClassComponent(element, this._context);
538+
this._updateClassComponent(elementType, element, this._context);
523539
} else {
524-
if (shouldConstruct(element.type)) {
525-
this._instance = new element.type(
540+
if (shouldConstruct(elementType)) {
541+
this._instance = new elementType(
526542
element.props,
527543
this._context,
528544
this._updater,
529545
);
530-
531-
if (typeof element.type.getDerivedStateFromProps === 'function') {
532-
const partialState = element.type.getDerivedStateFromProps.call(
546+
if (typeof elementType.getDerivedStateFromProps === 'function') {
547+
const partialState = elementType.getDerivedStateFromProps.call(
533548
null,
534549
element.props,
535550
this._instance.state,
@@ -543,39 +558,59 @@ class ReactShallowRenderer {
543558
}
544559
}
545560

546-
if (element.type.hasOwnProperty('contextTypes')) {
561+
if (elementType.contextTypes) {
547562
currentlyValidatingElement = element;
548-
549563
checkPropTypes(
550-
element.type.contextTypes,
564+
elementType.contextTypes,
551565
this._context,
552566
'context',
553-
getName(element.type, this._instance),
567+
getName(elementType, this._instance),
554568
getStackAddendum,
555569
);
556570

557571
currentlyValidatingElement = null;
558572
}
559573

560-
this._mountClassComponent(element, this._context);
574+
this._mountClassComponent(elementType, element, this._context);
561575
} else {
562-
const prevDispatcher = ReactCurrentDispatcher.current;
563-
ReactCurrentDispatcher.current = this._dispatcher;
564-
this._prepareToUseHooks(element.type);
565-
try {
566-
if (isForwardRef(element)) {
567-
this._rendered = element.type.render(element.props, element.ref);
568-
} else {
569-
this._rendered = element.type.call(
570-
undefined,
571-
element.props,
572-
this._context,
573-
);
576+
let shouldRender = true;
577+
if (
578+
isMemo(element.type) &&
579+
elementType === this._previousComponentIdentity &&
580+
previousElement !== null
581+
) {
582+
// This is a Memo component that is being re-rendered.
583+
const compare = element.type.compare || shallowEqual;
584+
if (compare(previousElement.props, element.props)) {
585+
shouldRender = false;
586+
}
587+
}
588+
if (shouldRender) {
589+
const prevDispatcher = ReactCurrentDispatcher.current;
590+
ReactCurrentDispatcher.current = this._dispatcher;
591+
this._prepareToUseHooks(elementType);
592+
try {
593+
// elementType could still be a ForwardRef if it was
594+
// nested inside Memo.
595+
if (elementType.$$typeof === ForwardRef) {
596+
invariant(
597+
typeof elementType.render === 'function',
598+
'forwardRef requires a render function but was given %s.',
599+
typeof elementType.render,
600+
);
601+
this._rendered = elementType.render.call(
602+
undefined,
603+
element.props,
604+
element.ref,
605+
);
606+
} else {
607+
this._rendered = elementType(element.props, this._context);
608+
}
609+
} finally {
610+
ReactCurrentDispatcher.current = prevDispatcher;
574611
}
575-
} finally {
576-
ReactCurrentDispatcher.current = prevDispatcher;
612+
this._finishHooks(element, context);
577613
}
578-
this._finishHooks(element, context);
579614
}
580615
}
581616

@@ -601,7 +636,11 @@ class ReactShallowRenderer {
601636
this._instance = null;
602637
}
603638

604-
_mountClassComponent(element: ReactElement, context: null | Object) {
639+
_mountClassComponent(
640+
elementType: Function,
641+
element: ReactElement,
642+
context: null | Object,
643+
) {
605644
this._instance.context = context;
606645
this._instance.props = element.props;
607646
this._instance.state = this._instance.state || null;
@@ -616,7 +655,7 @@ class ReactShallowRenderer {
616655
// In order to support react-lifecycles-compat polyfilled components,
617656
// Unsafe lifecycles should not be invoked for components using the new APIs.
618657
if (
619-
typeof element.type.getDerivedStateFromProps !== 'function' &&
658+
typeof elementType.getDerivedStateFromProps !== 'function' &&
620659
typeof this._instance.getSnapshotBeforeUpdate !== 'function'
621660
) {
622661
if (typeof this._instance.componentWillMount === 'function') {
@@ -638,8 +677,12 @@ class ReactShallowRenderer {
638677
// because DOM refs are not available.
639678
}
640679

641-
_updateClassComponent(element: ReactElement, context: null | Object) {
642-
const {props, type} = element;
680+
_updateClassComponent(
681+
elementType: Function,
682+
element: ReactElement,
683+
context: null | Object,
684+
) {
685+
const {props} = element;
643686

644687
const oldState = this._instance.state || emptyObject;
645688
const oldProps = this._instance.props;
@@ -648,7 +691,7 @@ class ReactShallowRenderer {
648691
// In order to support react-lifecycles-compat polyfilled components,
649692
// Unsafe lifecycles should not be invoked for components using the new APIs.
650693
if (
651-
typeof element.type.getDerivedStateFromProps !== 'function' &&
694+
typeof elementType.getDerivedStateFromProps !== 'function' &&
652695
typeof this._instance.getSnapshotBeforeUpdate !== 'function'
653696
) {
654697
if (typeof this._instance.componentWillReceiveProps === 'function') {
@@ -664,8 +707,8 @@ class ReactShallowRenderer {
664707

665708
// Read state after cWRP in case it calls setState
666709
let state = this._newState || oldState;
667-
if (typeof type.getDerivedStateFromProps === 'function') {
668-
const partialState = type.getDerivedStateFromProps.call(
710+
if (typeof elementType.getDerivedStateFromProps === 'function') {
711+
const partialState = elementType.getDerivedStateFromProps.call(
669712
null,
670713
props,
671714
state,
@@ -685,7 +728,10 @@ class ReactShallowRenderer {
685728
state,
686729
context,
687730
);
688-
} else if (type.prototype && type.prototype.isPureReactComponent) {
731+
} else if (
732+
elementType.prototype &&
733+
elementType.prototype.isPureReactComponent
734+
) {
689735
shouldUpdate =
690736
!shallowEqual(oldProps, props) || !shallowEqual(oldState, state);
691737
}
@@ -694,7 +740,7 @@ class ReactShallowRenderer {
694740
// In order to support react-lifecycles-compat polyfilled components,
695741
// Unsafe lifecycles should not be invoked for components using the new APIs.
696742
if (
697-
typeof element.type.getDerivedStateFromProps !== 'function' &&
743+
typeof elementType.getDerivedStateFromProps !== 'function' &&
698744
typeof this._instance.getSnapshotBeforeUpdate !== 'function'
699745
) {
700746
if (typeof this._instance.componentWillUpdate === 'function') {
@@ -729,7 +775,8 @@ function getDisplayName(element) {
729775
} else if (typeof element.type === 'string') {
730776
return element.type;
731777
} else {
732-
return element.type.displayName || element.type.name || 'Unknown';
778+
const elementType = isMemo(element.type) ? element.type.type : element.type;
779+
return elementType.displayName || elementType.name || 'Unknown';
733780
}
734781
}
735782

‎packages/react-test-renderer/src/__tests__/ReactShallowRenderer-test.js‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1454,4 +1454,115 @@ describe('ReactShallowRenderer', () => {
14541454
shallowRenderer.render(<Foo foo="bar" />);
14551455
expect(logs).toEqual([undefined]);
14561456
});
1457+
1458+
it('should handle memo', () => {
1459+
function Foo() {
1460+
return <div>foo</div>;
1461+
}
1462+
const MemoFoo = React.memo(Foo);
1463+
const shallowRenderer = createRenderer();
1464+
shallowRenderer.render(<MemoFoo />);
1465+
});
1466+
1467+
it('should enable React.memo to prevent a re-render', () => {
1468+
const logs = [];
1469+
const Foo = React.memo(({count}) => {
1470+
logs.push(`Foo: ${count}`);
1471+
return <div>{count}</div>;
1472+
});
1473+
const Bar = React.memo(({count}) => {
1474+
logs.push(`Bar: ${count}`);
1475+
return <div>{count}</div>;
1476+
});
1477+
const shallowRenderer = createRenderer();
1478+
shallowRenderer.render(<Foo count={1} />);
1479+
expect(logs).toEqual(['Foo: 1']);
1480+
logs.length = 0;
1481+
// Rendering the same element with the same props should be prevented
1482+
shallowRenderer.render(<Foo count={1} />);
1483+
expect(logs).toEqual([]);
1484+
// A different element with the same props should cause a re-render
1485+
shallowRenderer.render(<Bar count={1} />);
1486+
expect(logs).toEqual(['Bar: 1']);
1487+
});
1488+
1489+
it('should respect a custom comparison function with React.memo', () => {
1490+
let renderCount = 0;
1491+
function areEqual(props, nextProps) {
1492+
return props.foo === nextProps.foo;
1493+
}
1494+
const Foo = React.memo(({foo, bar}) => {
1495+
renderCount++;
1496+
return (
1497+
<div>
1498+
{foo} {bar}
1499+
</div>
1500+
);
1501+
}, areEqual);
1502+
1503+
const shallowRenderer = createRenderer();
1504+
shallowRenderer.render(<Foo foo={1} bar={1} />);
1505+
expect(renderCount).toBe(1);
1506+
// Change a prop that the comparison funciton ignores
1507+
shallowRenderer.render(<Foo foo={1} bar={2} />);
1508+
expect(renderCount).toBe(1);
1509+
shallowRenderer.render(<Foo foo={2} bar={2} />);
1510+
expect(renderCount).toBe(2);
1511+
});
1512+
1513+
it('should not call the comparison function with React.memo on the initial render', () => {
1514+
const areEqual = jest.fn(() => false);
1515+
const SomeComponent = React.memo(({foo}) => {
1516+
return <div>{foo}</div>;
1517+
}, areEqual);
1518+
const shallowRenderer = createRenderer();
1519+
shallowRenderer.render(<SomeComponent foo={1} />);
1520+
expect(areEqual).not.toHaveBeenCalled();
1521+
expect(shallowRenderer.getRenderOutput()).toEqual(<div>{1}</div>);
1522+
});
1523+
1524+
it('should handle memo(forwardRef())', () => {
1525+
const testRef = React.createRef();
1526+
const SomeComponent = React.forwardRef((props, ref) => {
1527+
expect(ref).toEqual(testRef);
1528+
return (
1529+
<div>
1530+
<span className="child1" />
1531+
<span className="child2" />
1532+
</div>
1533+
);
1534+
});
1535+
1536+
const SomeMemoComponent = React.memo(SomeComponent);
1537+
1538+
const shallowRenderer = createRenderer();
1539+
const result = shallowRenderer.render(<SomeMemoComponent ref={testRef} />);
1540+
1541+
expect(result.type).toBe('div');
1542+
expect(result.props.children).toEqual([
1543+
<span className="child1" />,
1544+
<span className="child2" />,
1545+
]);
1546+
});
1547+
1548+
it('should warn for forwardRef(memo())', () => {
1549+
const testRef = React.createRef();
1550+
const SomeMemoComponent = React.memo(({foo}) => {
1551+
return <div>{foo}</div>;
1552+
});
1553+
const shallowRenderer = createRenderer();
1554+
expect(() => {
1555+
expect(() => {
1556+
const SomeComponent = React.forwardRef(SomeMemoComponent);
1557+
shallowRenderer.render(<SomeComponent ref={testRef} />);
1558+
}).toWarnDev(
1559+
'Warning: forwardRef requires a render function but received ' +
1560+
'a `memo` component. Instead of forwardRef(memo(...)), use ' +
1561+
'memo(forwardRef(...))',
1562+
{withoutStack: true},
1563+
);
1564+
}).toThrowError(
1565+
'forwardRef requires a render function but was given object.',
1566+
);
1567+
});
14571568
});

0 commit comments

Comments
 (0)