Create a custom React hook to mimic class component's setState behaviour

I have been playing around with react hooks recently, and slowly wrapping my head around how they work. Hooks are a great way to create reusable functionality, but there are some instances where I think class components fared better. Let's talk about one such scenario, and also how to get around it when using hooks.

When working with react state in a class component, if we wanted to update a specific property on the state object, all we had to do was call the setState method and pass in an object containing only the updated property and value. The resultant state would be the previous state merged with the new property value.

It works a bit differently with hooks. When using the default useState hook, the new object passed in entirely replaces the previous state. So all the properties in the previous state are overwritten. So every time you want to update a state object, the onus is on the developer to keep track of the previous state and update the new state based on that. Here is how this behaviour manifests in code:

	const initialState: IState = {
	FirstName: 'Vardhaman',
	LastName: '',
	};
	//initialise useState hook with initialState
	const [state, setState] = React.useState<IState>(initialState);
	console.log({ state }); //state: {FirstName: "Vardhaman", LastName: ""}
	//Using the default useState hook replaces state object with the new object.
	setState({
	LastName: 'Deshpande'
	});
	console.log({ state }) //state: {LastName: "Deshpande"}

view raw useState.ts hosted with ❤ by GitHub

As you can see in the code the value of the FirstName property is lost when we update the LastName property. 

There is a way to set the state based on previous state by using the functional update pattern: https://reactjs.org/docs/hooks-reference.html#functional-updates

But that means that every time we used to just use this.setState in class components, we now have to use

setState(prevState => {

  return {...prevState, ...updatedValues};

});

This would be additional overhead for us devs which I wanted to check if we could avoid.

Fortunately with the reusable nature of hooks, we can simply create a custom hook which mimics the class components setState behaviour and merges the new state with the previous state. And we can use this custom hook any time we want to merge the state.

Here is how the previous code would look when using our custom hook:

	const initialState: IState = {
	FirstName: 'Vardhaman',
	LastName: '',
	};
	//initialise our custom useCustomState hook with initialState
	const [customState, setCustomState] = useCustomState<IState>(initialState);
	console.log({ customState }); //customState: {FirstName: "Vardhaman", LastName: ""}
	//Using our custom hook will merge the previous state object with new object mimicking a class components setState behavior
	setCustomState({
	LastName: 'Deshpande'
	});
	console.log({ customState }) //customState: {FirstName: "Vardhaman", LastName: "Deshpande"}

view raw useCustomState_consumer.ts hosted with ❤ by GitHub

The only change in code is that we have replaced react's useState with our custom useCustomState. This gives us back the ability to merge state objects instead of replacing them. (I have also changed the variable names from state and setState to customState and setCustomState respectively to make it explicit that we are using the custom hook. But you can just as well keep using the state and setState names)

So how does our custom hook work? It's built as a wrapper on top of the useState hook's functional update pattern. Any time a state object is passed to the setCustomState function, it internally uses the functional update pattern and merges the new state with the previous state. Let's have a look at the code:

	import { useCallback, useState } from 'react';
	/*
	Function returns
	1) A state object
	2) A function which accepts a partial state to merge
	*/
	export const useCustomState = <T extends object>(initialState: T): [T, (newPartialState: Partial<T>) => void] => {
	const [state, setState] = useState<T>(initialState);
	//function which accepts a partial state to merge
	const setCustomState = useCallback((newPartialState: T) => {
	try {
	setState((prevState): T => {
	return { ...prevState, ...newPartialState };
	});
	} catch (error) {
	console.error(error);
	}
	}, []);
	return [state, setCustomState];
	};

view raw useCustomState_basic.ts hosted with ❤ by GitHub

This automates the overhead of using the functional pattern. It is no longer the the developers responsibility, and instead, is done by the custom hook.

But wait, what if there is a scenario where a new state depends on the previous state? Our custom hook does not yet expose a way for the developer to update the state based on the previous state. Let's fix that. We can update our hook to add a method which accepts a function. This function will receive the previous state from react.  

	//A function which accepts an action to update the state based on previous state
	const setCustomStateDispatch = useCallback((newStateFunction: SetStateAction<T>) => {
	try {
	setState(newStateFunction);
	} catch (error) {
	console.error(error);
	}
	}, []);

view raw setCustomStateDispatch.ts hosted with ❤ by GitHub

And consuming the new hook can be done by passing in a function updater: 

	const initialState: IState = {
	FirstName: 'Vardhaman',
	LastName: 'Deshpande',
	FullName: ''
	;
	//initialise our custom useCustomState hook with initialState
	const [customState, setCustomState, setCustomStateDispatch] = useCustomState<IState>(initialState);
	console.log({ customState }); //customState: {FirstName: "Vardhaman", LastName: "Deshpande"}
	//Pass in a function which accepts a SetStateAction as a param. This will give us a hook into the previous State.
	setCustomStateDispatch((prevState: IState): IState => {
	return { ...prevState, FullName: `${prevState.FirstName} ${prevState.LastName}` };
	});
	console.log({ customState }); //customState: {FirstName: "Vardhaman", LastName: "Deshpande", FullName: "Vardhaman Deshpande"}

view raw customStateDispatch_consumer.ts hosted with ❤ by GitHub

Hope that helps! For the sake of completeness here is the full code for our custom hook:

	import { useCallback, useState, SetStateAction } from 'react';
	/*
	Function returns
	1) A state object
	2) A function which accepts a partial state to merge
	3) A function which accepts a SetStateAction to update the state based on previous state
	*/
	export const useCustomState = <T extends object>(initialState: T): [T, (newPartialState: Partial<T>) => void, (newStateFunction: SetStateAction<T>) => void] => {
	const [state, setState] = useState<T>(initialState);
	//function which accepts a partial state to merge
	const setCustomState = useCallback((newPartialState: T) => {
	try {
	setState((prevState): T => {
	return { ...prevState, ...newPartialState };
	});
	} catch (error) {
	console.error(error);
	}
	}, []);
	//A function which accepts an action to update the state based on previous state
	const setCustomStateDispatch = useCallback((newStateFunction: SetStateAction<T>) => {
	try {
	setState(newStateFunction);
	} catch (error) {
	console.error(error);
	}
	}, []);
	return [state, setCustomState, setCustomStateDispatch];
	};

view raw useCustomState.ts hosted with ❤ by GitHub

Full solution including the custom hook and the consuming code can be found here: https://github.com/vman/spfx-react-hooks-customstate

This post generated some good discussion on twitter with Yannick Plenevaux regarding the ideal cases when to use this approach as opposed to other approaches of state management like using the useState or useReducer hooks. Have a look here: https://twitter.com/yp_code/status/1265244244077416448