Zustand Automatically Generated Hooks & Get Selectors

zustand-selectors.ts

import { type StoreApi, useStore } from "zustand";

type WithHookSelectors<S> = S extends { getState: () => infer T }
  ? S & { use: { [K in keyof T]: () => T[K] } }
  : never;

export function createHookSelectors<S extends StoreApi<object>>(_store: S) {
  const store = _store as WithHookSelectors<typeof _store>;
  store.use = {};
  for (const key of Object.keys(store.getState())) {
    (store.use as any)[key] = () =>
      useStore(_store, (s) => s[key as keyof typeof s]);
  }

  return store;
}

type GetRecord<O> = {
  [K in keyof O]: (equalityFn?: (a: O[K], b: O[K]) => boolean) => O[K];
};

type WithGetSelectors<S> = S extends { getState: () => infer T }
  ? S & { get: GetRecord<T> }
  : never;

export function createGetSelectors<S extends StoreApi<object>>(_store: S) {
  const store = _store as WithGetSelectors<typeof _store>;
  store.get = {} as GetRecord<typeof store.getState>;

  for (const key of Object.keys(store.getState())) {
    (store.get as any)[key] = () =>
      store.getState()[key as keyof typeof store.getState];
  }

  return store;
}