From 1bfe0636643a519bcbc3e620d9e384b4870a5a8b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Oriol=20Colomer=20Aragon=C3=A9s?= <oriol.colomer@koobin.com>
Date: Thu, 24 Oct 2019 16:27:38 +0200
Subject: [PATCH] feat: add ensuredForwardRef and useEnsuredForwardedRef

---
 README.md                                     |  5 ++
 docs/useEnsuredForwardedRef.md                | 63 +++++++++++++++
 .../useEnsuredForwardedRef.story.tsx          | 79 +++++++++++++++++++
 src/__tests__/useEnsuredForwardedRef.test.tsx | 53 +++++++++++++
 src/index.ts                                  |  1 +
 src/useEnsuredForwardedRef.ts                 | 33 ++++++++
 6 files changed, 234 insertions(+)
 create mode 100644 docs/useEnsuredForwardedRef.md
 create mode 100644 src/__stories__/useEnsuredForwardedRef.story.tsx
 create mode 100644 src/__tests__/useEnsuredForwardedRef.test.tsx
 create mode 100644 src/useEnsuredForwardedRef.ts

diff --git a/README.md b/README.md
index 76fd92c187..d93ee9a387 100644
--- a/README.md
+++ b/README.md
@@ -138,6 +138,11 @@
   - [`useStateValidator`](./docs/useStateValidator.md) &mdash; tracks state of an object. [![][img-demo]](https://streamich.github.io/react-use/?path=/story/state-usestatevalidator--demo)
   - [`useMultiStateValidator`](./docs/useMultiStateValidator.md) &mdash; alike the `useStateValidator`, but tracks multiple states at a time. [![][img-demo]](https://streamich.github.io/react-use/?path=/story/state-usemultistatevalidator--demo)
   - [`useMediatedState`](./docs/useMediatedState.md) &mdash; like the regular `useState` but with mediation by custom function. [![][img-demo]](https://streamich.github.io/react-use/?path=/story/state-usemediatedstate--demo)
+    <br/>
+    <br/>
+- [**Miscellaneous**]()
+  - [`useEnsuredForwardedRef`](./docs/useEnsuredForwardedRef.md) and [`ensuredForwardRef`](./docs/useEnsuredForwardedRef.md) &mdash; use a React.forwardedRef safely. [![][img-demo]](https://streamich.github.io/react-use/?path=/story/state-useensuredforwardedref--demo)
+
 
 <br />
 <br />
diff --git a/docs/useEnsuredForwardedRef.md b/docs/useEnsuredForwardedRef.md
new file mode 100644
index 0000000000..b094a284bc
--- /dev/null
+++ b/docs/useEnsuredForwardedRef.md
@@ -0,0 +1,63 @@
+# `useEnsuredForwardedRef`
+
+React hook to use a ForwardedRef safely.
+
+In some scenarios, you may need to use a _ref_ from inside and outside a component. If that's the case, you should use `React.forwardRef` to pass it through the child component. This is useful when you only want to forward that _ref_ and expose an internal `HTMLelement` to a parent component, for example. However, if you need to manipulate that reference inside a child's lifecycle hook... things get complicated, since you can't always ensure that the _ref_ is being sent by the parent component and if it is not, you will get `undefined` instead of a valid _ref_.
+
+This hook is useful in this specific case, it will __ensure__ that you get a valid reference on the other side.
+
+## Usage
+
+```jsx
+import {ensuredForwardRef} from 'react-use';
+
+const Demo = () => {
+  return (
+    <Child />
+  );
+};
+
+const Child = ensuredForwardRef((props, ref) => {
+  useEffect(() => {
+    console.log(ref.current.getBoundingClientRect())
+  }, [])
+
+  return (
+    <div ref={ref} />
+  );
+});
+```
+
+## Alternative usage
+
+```jsx
+import {useEnsuredForwardedRef} from 'react-use';
+
+const Demo = () => {
+  return (
+    <Child />
+  );
+};
+
+const Child = React.forwardRef((props, ref) => {
+  // Here `ref` is undefined
+  const ensuredForwardRef = useEnsuredForwardedRef(ref);
+  // ensuredForwardRef will always be a valid reference.
+
+  useEffect(() => {
+    console.log(ensuredForwardRef.current.getBoundingClientRect())
+  }, [])
+
+  return (
+    <div ref={ensuredForwardRef} />
+  );
+});
+```
+
+## Reference
+
+```ts
+ensuredForwardRef<T, P = {}>(Component: RefForwardingComponent<T, P>): ForwardRefExoticComponent<PropsWithoutRef<P> & RefAttributes<T>>;
+
+useEnsuredForwardedRef<T>(ref: React.MutableRefObject<T>): React.MutableRefObject<T>;
+```
diff --git a/src/__stories__/useEnsuredForwardedRef.story.tsx b/src/__stories__/useEnsuredForwardedRef.story.tsx
new file mode 100644
index 0000000000..b2623f2fc4
--- /dev/null
+++ b/src/__stories__/useEnsuredForwardedRef.story.tsx
@@ -0,0 +1,79 @@
+import { storiesOf } from '@storybook/react';
+import React, { forwardRef, useRef, useState, useEffect, MutableRefObject } from 'react';
+import { useEnsuredForwardedRef } from '..';
+import ShowDocs from './util/ShowDocs';
+
+import { boolean, withKnobs } from '@storybook/addon-knobs';
+
+const INITIAL_SIZE = {
+  width: null,
+  height: null,
+};
+
+const Demo = ({ activeForwardRef }) => {
+  const ref = useRef(null);
+
+  const [size, setSize] = useState(INITIAL_SIZE);
+
+  useEffect(() => {
+    handleClick();
+  }, [activeForwardRef]);
+
+  const handleClick = () => {
+    if (activeForwardRef) {
+      const { width, height } = ref.current.getBoundingClientRect();
+      setSize({
+        width,
+        height,
+      });
+    } else {
+      setSize(INITIAL_SIZE);
+    }
+  };
+
+  return (
+    <>
+      <button onClick={handleClick} disabled={!activeForwardRef}>
+        {activeForwardRef ? 'Update parent component' : 'forwardRef value is undefined'}
+      </button>
+      <div>Parent component using external ref: (textarea size)</div>
+      <pre>{JSON.stringify(size, null, 2)}</pre>
+      <Child ref={activeForwardRef ? ref : undefined} />
+    </>
+  );
+};
+
+const Child = forwardRef(({}, ref: MutableRefObject<HTMLTextAreaElement>) => {
+  const ensuredForwardRef = useEnsuredForwardedRef(ref);
+
+  const [size, setSize] = useState(INITIAL_SIZE);
+
+  useEffect(() => {
+    handleMouseUp();
+  }, []);
+
+  const handleMouseUp = () => {
+    const { width, height } = ensuredForwardRef.current.getBoundingClientRect();
+    setSize({
+      width,
+      height,
+    });
+  };
+
+  return (
+    <>
+      <div>Child forwardRef component using forwardRef: (textarea size)</div>
+      <pre>{JSON.stringify(size, null, 2)}</pre>
+      <div>You can resize this textarea:</div>
+      <textarea ref={ensuredForwardRef} onMouseUp={handleMouseUp} />
+    </>
+  );
+});
+
+storiesOf('Miscellaneous|useEnsuredForwardedRef', module)
+  .addDecorator(withKnobs)
+  .add('Docs', () => <ShowDocs md={require('../../docs/useEnsuredForwardedRef.md')} />)
+  .add('Demo', () => {
+    const activeForwardRef = boolean('activeForwardRef', true);
+    return <Demo activeForwardRef={activeForwardRef} />;
+  });
diff --git a/src/__tests__/useEnsuredForwardedRef.test.tsx b/src/__tests__/useEnsuredForwardedRef.test.tsx
new file mode 100644
index 0000000000..7f161d1a2e
--- /dev/null
+++ b/src/__tests__/useEnsuredForwardedRef.test.tsx
@@ -0,0 +1,53 @@
+import React, { useRef } from 'react';
+import ReactDOM from 'react-dom';
+import { renderHook } from '@testing-library/react-hooks';
+import TestUtils from 'react-dom/test-utils';
+import { useEnsuredForwardedRef } from '..';
+
+let container: HTMLDivElement;
+
+beforeEach(() => {
+  container = document.createElement('div');
+  document.body.appendChild(container);
+});
+
+afterEach(() => {
+  document.body.removeChild(container);
+  container = null;
+});
+
+test('should return a valid ref with existing forwardedRef', () => {
+  const { result } = renderHook(() => {
+    const ref = useRef(null);
+    const ensuredRef = useEnsuredForwardedRef(ref);
+
+    TestUtils.act(() => {
+      ReactDOM.render(<div ref={ensuredRef} />, container);
+    });
+
+    return {
+      initialRef: ref,
+      ensuredForwardedRef: ensuredRef,
+    };
+  });
+
+  const { initialRef, ensuredForwardedRef } = result.current;
+
+  expect(ensuredForwardedRef).toStrictEqual(initialRef);
+});
+
+test('should return a valid ref when the forwarded ref is undefined', () => {
+  const { result } = renderHook(() => {
+    const ref = useEnsuredForwardedRef<HTMLDivElement>(undefined);
+
+    TestUtils.act(() => {
+      ReactDOM.render(<div id="test_id" ref={ref} />, container);
+    });
+
+    return { ensuredRef: ref };
+  });
+
+  const { ensuredRef } = result.current;
+
+  expect(ensuredRef.current.id).toBe('test_id');
+});
diff --git a/src/index.ts b/src/index.ts
index 9c660296a2..fb77bf8407 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -17,6 +17,7 @@ export { default as useDefault } from './useDefault';
 export { default as useDrop } from './useDrop';
 export { default as useDropArea } from './useDropArea';
 export { default as useEffectOnce } from './useEffectOnce';
+export { default as useEnsuredForwardedRef, ensuredForwardRef } from './useEnsuredForwardedRef';
 export { default as useEvent } from './useEvent';
 export { default as useFavicon } from './useFavicon';
 export { default as useFullscreen } from './useFullscreen';
diff --git a/src/useEnsuredForwardedRef.ts b/src/useEnsuredForwardedRef.ts
new file mode 100644
index 0000000000..d7451e55a4
--- /dev/null
+++ b/src/useEnsuredForwardedRef.ts
@@ -0,0 +1,33 @@
+import {
+  forwardRef,
+  useRef,
+  useEffect,
+  MutableRefObject,
+  ForwardRefExoticComponent,
+  PropsWithoutRef,
+  RefAttributes,
+  RefForwardingComponent,
+  PropsWithChildren,
+} from 'react';
+
+export default function useEnsuredForwardedRef<T>(forwardedRef: MutableRefObject<T>): MutableRefObject<T> {
+  const ensuredRef = useRef(forwardedRef && forwardedRef.current);
+
+  useEffect(() => {
+    if (!forwardedRef) {
+      return;
+    }
+    forwardedRef.current = ensuredRef.current;
+  }, [forwardedRef]);
+
+  return ensuredRef;
+}
+
+export function ensuredForwardRef<T, P = {}>(
+  Component: RefForwardingComponent<T, P>
+): ForwardRefExoticComponent<PropsWithoutRef<P> & RefAttributes<T>> {
+  return forwardRef((props: PropsWithChildren<P>, ref) => {
+    const ensuredRef = useEnsuredForwardedRef(ref as MutableRefObject<T>);
+    return Component(props, ensuredRef);
+  });
+}