A React Hook to set persisted state, similar to useState

Example.tsx

import React from 'react';
import ReactDOM from 'react-dom';
import usePersistedState from './usePersistedState';

function App() {
  const [count, setCount, unsetCount] = usePersistedState(
    0,
    'count',
    sessionStorage
  );

  return (
    <div className="App">
      <p>You have clicked {count} times</p>
      <button onClick={() => setCount(prevCount => prevCount + 1)}>
        Click
      </button>
      <button onClick={() => unsetCount()}>Unset</button>
    </div>
  );
}

const rootElement = document.getElementById('root');
ReactDOM.render(<App />, rootElement);

usePersistedState.ts

import { useState } from 'react';

export default function usePersistedState<T>(
  initialState: T,
  key: string,
  storage: Storage = localStorage
): [T, (updater: T | ((previousState: T) => T)) => void, () => void] {
  let cachedValue = localStorage.getItem(key);
  let value: T;

  if (!cachedValue) {
    value = initialState;
  } else {
    value = JSON.parse(cachedValue);
  }

  const [state, setState] = useState(value);

  const setPersistedState = (nextState: T | ((previousState: T) => T)) => {
    if (typeof nextState === 'function') {
      return setState(previousState => {
        const newState = nextState(previousState);
        localStorage.setItem(key, JSON.stringify(newState));
        return newState;
      });
    }

    localStorage.setItem(key, JSON.stringify(nextState));
    setState(nextState);
  };

  const unsetPersistedState = () => {
    localStorage.removeItem(key);
  };

  return [state, setPersistedState, unsetPersistedState];
}