Zustand TypeScript Gotchas

There’s a couple gotchas I often face when using Zustand in my projects. I’m documenting them here for future Matt’s sake.

Usage with React Context

The simplest usage of Zustand is the create function that creates a global state hook. In some cases it’s preferable to scope this state explicity with React Context.

Their README shows how to achieve this with top level store (source). I prefer putting it in useState so that each provider owns a separate copy of the store, making providers more reusable.

      type MyStore = {
	count: number
	increment: () => void
}

function createMyStore() {
	return createStore<MyStore>((set) => ({
		count: 0,
		increment: () => set((state) => ({ count: state.count + 1 })),
	}))
}

type MyStoreApi = ReturnType<typeof createMyStore>

const MyStoreContext = createContext<MyStoreApi | null>(null)

export function MyStoreProvider({ children }: { children: React.ReactNode }) {
	const [store] = useState(createMyStore)
	return (
		<MyStoreContext.Provider value={store}>{children}</MyStoreContext.Provider>
	)
}

    

Usually with vanilla store you need to hand it to Zustand’s useStore to consume it. Add in React’s useContext and now you have to call two hooks each time you read state.

      import { useContext } from "react"
import { useStore } from "zustand"

function MyStoreConsumer() {
	/**
	 * I don't want to do this everywhere I read state
	 */
	const store = useContext(MyStoreContext)
	const count = useStore(store, (state) => state.count)

	return <p>{count}</p>
}

    

In order to create a bounded hook for this vanilla store, you can take advantage of function overloads and a non-null assertion to please the compiler and linter. I thought I was pretty clever for coming up with these types, but while writing this post I realized this exact pattern is already documented in Zustand’s TypeScript Guide—d’oh!

      // The first two lines are "Overload Signatures", the third is the "Implementation Signature"
export function useMyStore(): MyStore
export function useMyStore<T>(selector: (state: MyStore) => T): T
export function useMyStore<T>(selector?: (state: MyStore) => T) {
	const store = useContext(MyStoreContext)
	if (store === null) throw new Error("Missing <MyStoreProvider>")

	// Non-null assertion silences `useStore` complaint about second argument being `undefined` (not actually a runtime issue)
	return useStore(store, selector!)
}

    

Now we can read our state just like we would with the hook produced from create.

      function MyStoreConsumer() {
	const state = useMyStore() // get entire store
	// ^? const store: MyStore
	const count = useMyStore((state) => state.count) // select specific item
	// ^? const count: number

	return <p>{count}</p>
}

    

You may also want a hook that simply returns the entire store object from context, so you call call its imperative methods like .getState, .setState and .subscribe

      export function useMyStoreContext() {
	const store = useContext(MyStoreContext)
	if (store === null) throw new Error("Missing <MyStoreProvider>")

	return store
}

function MyStoreConsumer() {
	const store = useMyStoreContext()

	useEffect(() => {
		return store.subscribe(console.log)
	}, [])

	return null
}

    

Usage with Middleware

When using TypeScript, you should use the curried version of create/createStore. This means calling create<State>()(...) instead of create(). Note the extra pair of parenthesis. The reasoning is explained in more detail in their TypeScript guide.

I didn’t notice this at first because the non-curried version works too, until you try using middleware and face some dense type errors:

      import { create } from "zustand"
import { subscribeWithSelector } from "zustand/middleware"

const useStore = create<MyStore>(
	subscribeWithSelector((set) => ({
		count: 0,
		increment: () => set((state) => ({ count: state.count + 1 })),
	})),
)
// Error: Argument of type 'StateCreator<MyStore, [], [["zustand/subscribeWithSelector", never]]>' is not assignable to parameter of type 'StateCreator<MyStore, [], [], MyStore>'.

    

Simply add in those extra parenthesis and you’ll be good to go:

      import { create } from "zustand"
import { subscribeWithSelector } from "zustand/middleware"

const useStore = create<MyStore>()(
	subscribeWithSelector((set) => ({
		count: 0,
		increment: () => set((state) => ({ count: state.count + 1 })),
	})),
)

    

Complete Example

      import { useState, createContext, useContext, useEffect } from "react"
import { createStore, useStore, type StoreApi } from "zustand"
import { subscribeWithSelector } from "zustand/middleware"

type MyStore = {
	count: number
	increment: () => void
}

function createMyStore() {
	return createStore<MyStore>()(
		subscribeWithSelector((set) => ({
			count: 0,
			increment: () => set((state) => ({ count: state.count + 1 })),
		})),
	)
}

type MyStoreApi = ReturnType<typeof createMyStore>

const MyStoreContext = createContext<MyStoreApi | null>(null)

export function MyStoreProvider({ children }: { children: React.ReactNode }) {
	const [store] = useState(createMyStore)
	return (
		<MyStoreContext.Provider value={store}>{children}</MyStoreContext.Provider>
	)
}

export function useMyStore(): MyStore
export function useMyStore<T>(selector: (state: MyStore) => T): T
export function useMyStore<T>(selector?: (state: MyStore) => T) {
	const store = useContext(MyStoreContext)
	if (store === null) throw new Error("Missing <MyStoreProvider>")

	return useStore(store, selector!)
}

export function useMyStoreContext() {
	const store = useContext(MyStoreContext)
	if (store === null) throw new Error("Missing <MyStoreProvider>")

	return store
}

export function MyStoreConsumer() {
	const state = useMyStore() // get entire store
	const count = useMyStore((state) => state.count) // select specific item

	const store = useMyStoreContext()
	useEffect(() => {
		return store.subscribe(
			(state) => state.count,
			(count) => {
				console.log("count:", count) // subscribe with selector
			},
		)
	}, [])

	return <p>{count}</p>
}