From ed9097f92c9079f77f602343c3291ff6bc94c597 Mon Sep 17 00:00:00 2001
From: Jason Quense <monastic.panic@gmail.com>
Date: Tue, 22 Oct 2019 14:19:53 -0400
Subject: [PATCH] feat: add useSafeState

---
 src/useSafeState.ts        | 41 ++++++++++++++++++++++++++
 src/useSet.ts              |  2 +-
 test/useSafeState.test.tsx | 59 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 101 insertions(+), 1 deletion(-)
 create mode 100644 src/useSafeState.ts
 create mode 100644 test/useSafeState.test.tsx

diff --git a/src/useSafeState.ts b/src/useSafeState.ts
new file mode 100644
index 0000000..531ebd3
--- /dev/null
+++ b/src/useSafeState.ts
@@ -0,0 +1,41 @@
+import { Dispatch, SetStateAction, useCallback } from 'react'
+import useMounted from './useMounted'
+import { AsyncSetState } from './useStateAsync'
+
+type StateSetter<TState> = Dispatch<SetStateAction<TState>>
+
+/**
+ * `useSafeState` takes the return value of a `useState` hook and wraps the
+ * setter to prevent updates onces the component has unmounted. Can used
+ * with `useMergeState` and `useStateAsync` as well
+ *
+ * @param state The return value of a useStateHook
+ *
+ * ```ts
+ * const [show, setShow] = useSafeState(useState(true));
+ * ```
+ */
+function useSafeState<TState>(
+  state: [TState, AsyncSetState<TState>],
+): [TState, (stateUpdate: React.SetStateAction<TState>) => Promise<void>]
+function useSafeState<TState>(
+  state: [TState, StateSetter<TState>],
+): [TState, StateSetter<TState>]
+function useSafeState<TState>(
+  state: [TState, StateSetter<TState> | AsyncSetState<TState>],
+): [TState, StateSetter<TState> | AsyncSetState<TState>] {
+  const isMounted = useMounted()
+
+  return [
+    state[0],
+    useCallback(
+      (nextState: SetStateAction<TState>) => {
+        if (!isMounted()) return
+        return state[1](nextState)
+      },
+      [isMounted, state[1]],
+    ),
+  ]
+}
+
+export default useSafeState
diff --git a/src/useSet.ts b/src/useSet.ts
index 2e779ad..1d8f9ae 100644
--- a/src/useSet.ts
+++ b/src/useSet.ts
@@ -33,7 +33,7 @@ export class ObservableSet<V> extends Set<V> {
 /**
  * Create and return a [Set](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Set) that triggers rerenders when it's updated.
  *
- * ```tsx
+ * ```ts
  * const ids = useSet<number>([1,2,3,4]);
  *
  * return (
diff --git a/test/useSafeState.test.tsx b/test/useSafeState.test.tsx
new file mode 100644
index 0000000..662eeca
--- /dev/null
+++ b/test/useSafeState.test.tsx
@@ -0,0 +1,59 @@
+import { mount } from 'enzyme'
+import React, { useEffect, useState } from 'react'
+import { act } from 'react-dom/test-utils'
+import useSafeState from '../src/useSafeState'
+import useStateAsync from '../src/useStateAsync'
+
+describe('useSafeState', () => {
+  it('should work transparently', () => {
+    let state
+
+    function Wrapper() {
+      state = useSafeState(useState(false))
+      return null
+    }
+
+    const wrapper = mount(<Wrapper />)
+
+    expect(state[0]).toEqual(false)
+
+    act(() => {
+      state[1](true)
+    })
+    expect(state[0]).toEqual(true)
+
+    wrapper.unmount()
+
+    act(() => {
+      state[1](false)
+    })
+    expect(state[0]).toEqual(true)
+  })
+
+  it('should work with async setState', async () => {
+    let state
+
+    function Wrapper() {
+      state = useSafeState(useStateAsync(false))
+      return null
+    }
+
+    const wrapper = mount(<Wrapper />)
+
+    expect(state[0]).toEqual(false)
+
+    await act(async () => {
+      await state[1](true)
+    })
+
+    expect(state[0]).toEqual(true)
+
+    wrapper.unmount()
+
+    await act(async () => {
+      await state[1](true)
+    })
+
+    expect(state[0]).toEqual(true)
+  })
+})