Monday, January 15, 2024

Zustand, devtools and immer wrapper, pt. 2

On further researching of my previous idea of creating a Zustand + Immer wrapper, I found a very interesting technique straight from Zustand docs, where all the methods are simply free functions declared outside the store, called no store actions. I was already using this to declare private store methods, but it never occurred me to use it for all methods.

Since it officially has no downsides, I applied this idea to append the actions directly to the hook itself, so we don’t need to import each free function. The problem is that the hook has a couple methods itself, plus the inherent ones to any function object. So the result is very polluted. So I decided to adopt a convention of prefixing every action with $.

But conventions are easy to break. So I implemented a reduce to prefix each method name with #. But then TypeScript could not infer the type anymore, so I had to dig and type it manually with key remapping, and finally I got this implementation:

import {create, StateCreator} from 'zustand';
import {devtools} from 'zustand/middleware';
import {immer} from 'zustand/middleware/immer';
	
export function createStore<T extends object, U extends object>(
	name: string,
	initialState: T,
	actions: StateCreator<
		T,
		[['zustand/devtools', never], ['zustand/immer', never]],
		[['zustand/immer', never], ['zustand/devtools', never]],
		U
	>,
) {
	const store = create(
		devtools(
			immer(() => initialState),
			{name: name, serialize: true},
		),
	);

	type WithPrefix<T> = {
		[K in keyof T as K extends string ? `$${K}` : never]: T[K];
	}
	const actionsOrig = actions(store.setState, store.getState, store);
	const actionsPrefixed = Object.entries(actionsOrig).reduce(
		(accum, [k, v]) => ({...accum, ['$' + k]: v}),
		{} as WithPrefix<typeof actionsOrig>,
	);

	return Object.assign(store, actionsPrefixed);
}

This creator allows us to call actions directly, without the hook. They’ll be automatically prefixed with $:

const useFoo = createStore('Foo', {
	name: '',
}, (set, get) => ({
	setName(name: string): void {
		set(state => {
			state.name = name;
		});
	},
}));
	
function App() {
	const name = useFoo(s => s.name);

	return <div>
		<span>{name}</span>
		<button onClick={() => useFoo.$setName('Joe')}</button>
			Change
		</button>
	</div>;
}

That $ prefix gives me bad Vue vibes. Let’s see if I can get along with it, because I was able to develop a DX which is exactly what I was looking for.

No comments: