Recreating this.setState() with React Hooks

React Hooks are far from new, so there's no point in a post that sounds like a broken record. Instead, I'd like to take some of your time to introduce a small recipe that I hacked together using hooks.

In this post I will explain how I re-created the traditional this.setState() method of a React class component using hooks. We will create a custom hook useSetState that will return a tuple [state, setState] that behaves like this.state and this.setState respectively.

The useState hook

Creating the component state through hooks is done via the useState hook. Thus the initial part was just returning the values received from useState from the hook. The only thing we need to take care of is partial state updates. If you recall, this.setState merges its argument object with the current state object to obtain the new state, which is in contrast to the updater function returned by useState which completely replaces the corresponding state slice with whatever argument given to it. So, the code at this point looks like:

const useSetState = (initState) => {
  const [_state, _setState] = useState(initState);

  const setState = (update) => {
    const newState = {
      ..._state,
      ...update,
    };

    _setState(newState);
  };

  return [_state, setState];
};

The updater argument to this.setState

Even though most developers update the state using an object, there are cases where you need a function to update the state (eg: when the current state depends on the previous state.) In fact, my most popular answer on Stack Overflow is one that suggests the use of the "updater argument" to this.setState.
In order to support this argument type as well, we need to have a way to update state based on previous state. Sort of like, updater(prevState) => nextState. Wait, isn't that a reducer??

So now, let's ditch useState and use useReducer instead, with the same functionality. We'll support the updater argument, but not yet.

import React, { useReducer } from 'react';

const PATCH = '@action_types/PATCH';

const reducer = (state, action) => {
  if ( action.type === PATCH ) {
    return {
      ...state,
      ...action.payload,
    };
  }
};

const useSetState = (initState) => {
  const [_state, _dispatch] = useReducer(reducer, initState);
  const _patchState = update => _dispatch({ type: PATCH, payload: update });

  const setState = (update) => {
    const newState = {
      ..._state,
      ...update,
    };

    _patchState(newState);
  };

  return [_state, setState];
};

Now we'll add the updater argument:

import { useReducer } from 'react';

const PATCH = '@action_types/PATCH';
const DERIVE = '@action_types/DERIVE';

const reducer = (state, action) => {
  switch ( action.type ) {
    case PATCH:
      return {
        ...state,
        ...action.payload,
      };
    case DERIVE:
      return {
        ...state,
        ...action.updater(state),
      };
    default: console.error(`Unexpected action type: ${action.type}`); return state;
  }
};

const useSetState = (initState) => {
  const [_state, _dispatch] = useReducer(reducer, initState);
  const _patchState = update => _dispatch({ type: PATCH, payload: update });
  const _deriveState = updater => _dispatch({ type: DERIVE, updater });

  const setState = (arg) => {
    if ( typeof arg === 'function' ) {
      _deriveState(arg);
    } else {
      _patchState(arg);
    }
  };

  return [_state, setState];
};

export default useSetState;

We can see how 2 action types DERIVE and PATCH are used to represent the 2 types of changes that may happen to the state.

The Last Piece

It so happens that this.setState supports a second argument. From the React docs:

The second parameter to setState() is an optional callback function that will be executed once setState is completed and the component is re-rendered. Generally we recommend using componentDidUpdate() for such logic instead.

And use componentDidUpdate is what we're gonna do. Or at least the hooks equivalent of it. If you know how useEffect works, running a piece of code every time some data changes is trivial. If not, I recommend reading the useEffect doc.

So, yes, we're going to run the second argument to our setState function after the state changes. But how do we store the function somewhere such that its value is not lost/reset across renders?
Enter useRef. As soon as setState is called, we save the second argument in a ref object. Then in the useEffect callback, we

1. Invoke the function stored in the ref object, and
2. Clear the ref object ( or set to no-op )

With this, we're done, and the final code ( after adding some type checks ) looks like this:

import { useReducer, useEffect, useRef } from 'react';

const PATCH = '@action_types/PATCH';
const DERIVE = '@action_types/DERIVE';

const noop = () => {};

const isObject = (arg) => {
  return arg === Object(arg) && !Array.isArray(arg);
};

const reducer = (state, action) => {
  switch ( action.type ) {
    case PATCH:
      return {
        ...state,
        ...action.payload,
      };
    case DERIVE:
      return {
        ...state,
        ...action.updater(state),
      };
    default: console.error(`Unexpected action type: ${action.type}`); return state;
  }
};

const useSetState = (initState) => {
  const [_state, _dispatch] = useReducer(reducer, initState);

  const _patchState = update => _dispatch({ type: PATCH, payload: update });
  const _deriveState = updater => _dispatch({ type: DERIVE, updater });

  const _setStateCallback = useRef();

  useEffect(() => {
    if ( typeof _setStateCallback.current === 'function' ) {
      _setStateCallback.current();
    }
    _setStateCallback.current = noop;
  }, [_state]);

  const setState = (arg, callback = noop) => {
    _setStateCallback.current = callback;
    if ( typeof arg === 'function' ) {
      _deriveState(arg);
    } else if ( isObject(arg) ) {
      _patchState(arg);
    } else {
      throw Error(
        'Invalid argument type passed to setState. Argument must either be a plain object or' +
        'an updater function.'
      );
    }
  };

  return [_state, setState];
};

export default useSetState;

Conclusion

Like I've written before, the ability to create custom hooks is one of the biggest boons that React Hooks bring with them. As long as you are using functional components, React hooks is one of the best ways to:

1. Create custom recipes/functionalities like this one, and
2. Create reusable logic that can be shared across components

The above hook is published as an npm package called @danedavid/usesetstate. Initially I didn't want to publish this package as there were other similar packages in npm. But in case you want to test it, there it is. The code can be found here and the hook can be found in action here.

---

If you liked this post or have any questions, please do comment below, or shoot me an email at dndavid102[at]gmail[dot]com. You may also follow me on Twitter.