Support ref cleanup function for imperative handle refs (#28910)
Support ref cleanup function for imperative handle refs
Jan Kassens committed
Apr 25, 2024 at 12:51 UTC
ed71a3ad2965617c27c6e7ca7577f15b8ca4152c
2 files changed
+103
-17
packages/react-dom/src/__tests__/refs-test.js
+96
-15
@@ -9,9 +9,9 @@
9
10
'use strict';
11
12
-let React = require('react');
13
-let ReactDOMClient = require('react-dom/client');
14
-let act = require('internal-test-utils').act;
12
+const React = require('react');
13
+const ReactDOMClient = require('react-dom/client');
14
+const act = require('internal-test-utils').act;
15
16
// This is testing if string refs are deleted from `instance.refs`
17
// Once support for string refs is removed, this test can be removed.
@@ -19,13 +19,6 @@ let act = require('internal-test-utils').act;
19
describe('reactiverefs', () => {
20
let container;
21
22
- beforeEach(() => {
23
- jest.resetModules();
24
- React = require('react');
25
- ReactDOMClient = require('react-dom/client');
26
- act = require('internal-test-utils').act;
27
- });
28
-
22
afterEach(() => {
23
if (container) {
24
document.body.removeChild(container);
@@ -199,11 +192,6 @@ describe('reactiverefs', () => {
192
describe('ref swapping', () => {
193
let RefHopsAround;
194
beforeEach(() => {
202
- jest.resetModules();
203
- React = require('react');
204
- ReactDOMClient = require('react-dom/client');
205
- act = require('internal-test-utils').act;
206
-
195
RefHopsAround = class extends React.Component {
196
container = null;
197
state = {count: 0};
@@ -804,3 +792,96 @@ describe('refs return clean up function', () => {
792
expect(nullHandler).toHaveBeenCalledTimes(0);
793
});
794
});
795
+
796
+describe('useImerativeHandle refs', () => {
797
+ const ImperativeHandleComponent = React.forwardRef(({name}, ref) => {
798
+ React.useImperativeHandle(
799
+ ref,
800
+ () => ({
801
+ greet() {
802
+ return `Hello ${name}`;
803
+ },
804
+ }),
805
+ [name],
806
+ );
807
+ return null;
808
+ });
809
+
810
+ it('should work with object style refs', async () => {
811
+ const container = document.createElement('div');
812
+ const root = ReactDOMClient.createRoot(container);
813
+ const ref = React.createRef();
814
+
815
+ await act(async () => {
816
+ root.render(<ImperativeHandleComponent name="Alice" ref={ref} />);
817
+ });
818
+ expect(ref.current.greet()).toBe('Hello Alice');
819
+ await act(() => {
820
+ root.render(null);
821
+ });
822
+ expect(ref.current).toBe(null);
823
+ });
824
+
825
+ it('should work with callback style refs', async () => {
826
+ const container = document.createElement('div');
827
+ const root = ReactDOMClient.createRoot(container);
828
+ let current = null;
829
+
830
+ await act(async () => {
831
+ root.render(
832
+ <ImperativeHandleComponent
833
+ name="Alice"
834
+ ref={r => {
835
+ current = r;
836
+ }}
837
+ />,
838
+ );
839
+ });
840
+ expect(current.greet()).toBe('Hello Alice');
841
+ await act(() => {
842
+ root.render(null);
843
+ });
844
+ expect(current).toBe(null);
845
+ });
846
+
847
+ it('should work with callback style refs with cleanup function', async () => {
848
+ const container = document.createElement('div');
849
+ const root = ReactDOMClient.createRoot(container);
850
+
851
+ let cleanupCalls = 0;
852
+ let createCalls = 0;
853
+ let current = null;
854
+
855
+ const ref = r => {
856
+ current = r;
857
+ createCalls++;
858
+ return () => {
859
+ current = null;
860
+ cleanupCalls++;
861
+ };
862
+ };
863
+
864
+ await act(async () => {
865
+ root.render(<ImperativeHandleComponent name="Alice" ref={ref} />);
866
+ });
867
+ expect(current.greet()).toBe('Hello Alice');
868
+ expect(createCalls).toBe(1);
869
+ expect(cleanupCalls).toBe(0);
870
+
871
+ // update a dep should recreate the ref
872
+ await act(async () => {
873
+ root.render(<ImperativeHandleComponent name="Bob" ref={ref} />);
874
+ });
875
+ expect(current.greet()).toBe('Hello Bob');
876
+ expect(createCalls).toBe(2);
877
+ expect(cleanupCalls).toBe(1);
878
+
879
+ // unmounting should call cleanup
880
+ await act(() => {
881
+ root.render(null);
882
+ });
883
+ expect(current).toBe(null);
884
+ expect(createCalls).toBe(2);
885
+ expect(cleanupCalls).toBe(2);
886
+ });
887
+});
packages/react-reconciler/src/ReactFiberHooks.js
+7
-2
@@ -2564,9 +2564,14 @@ function imperativeHandleEffect<T>(
2564
if (typeof ref === 'function') {
2565
const refCallback = ref;
2566
const inst = create();
2567
- refCallback(inst);
2567
+ const refCleanup = refCallback(inst);
2568
return () => {
2569
- refCallback(null);
2569
+ if (typeof refCleanup === 'function') {
2570
+ // $FlowFixMe[incompatible-use] we need to assume no parameters
2571
+ refCleanup();
2572
+ } else {
2573
+ refCallback(null);
2574
+ }
2575
};
2576
} else if (ref !== null && ref !== undefined) {
2577
const refObject = ref;