@assistant-ui/react
Version:
React components for AI chat.
75 lines (68 loc) • 2.39 kB
text/typescript
import { UseBoundStore } from "zustand";
import { ReadonlyStore } from "../../ReadonlyStore";
/**
* Creates hooks for accessing a store within a context.
* @param contextHook - The hook to access the context.
* @param contextKey - The key of the store in the context.
* @returns An object containing the hooks: `use...` and `use...Store`.
*/
export function createContextStoreHook<T, K extends keyof T & string>(
contextHook: (options?: { optional?: boolean }) => T | null,
contextKey: K,
) {
type StoreType = T[K];
type StateType = StoreType extends ReadonlyStore<infer S> ? S : never;
// Define useStoreStoreHook with overloads
function useStoreStoreHook(): ReadonlyStore<StateType>;
function useStoreStoreHook(options: {
optional: true;
}): ReadonlyStore<StateType> | null;
function useStoreStoreHook(options?: {
optional?: boolean;
}): ReadonlyStore<StateType> | null {
const context = contextHook(options);
if (!context) return null;
return context[contextKey] as ReadonlyStore<StateType>;
}
// Define useStoreHook with overloads
function useStoreHook(): StateType;
function useStoreHook<TSelected>(
selector: (state: StateType) => TSelected,
): TSelected;
function useStoreHook(options: { optional: true }): StateType | null;
function useStoreHook<TSelected>(options: {
optional: true;
selector?: (state: StateType) => TSelected;
}): TSelected | null;
function useStoreHook<TSelected>(
param?:
| ((state: StateType) => TSelected)
| {
optional?: boolean;
selector?: (state: StateType) => TSelected;
},
): TSelected | StateType | null {
let optional = false;
let selector: ((state: StateType) => TSelected) | undefined;
if (typeof param === "function") {
selector = param;
} else if (param && typeof param === "object") {
optional = !!param.optional;
selector = param.selector;
}
const store = useStoreStoreHook({
optional,
} as any) as UseBoundStore<ReadonlyStore<StateType>>;
if (!store) return null;
return selector ? store(selector) : store();
}
// Return an object with keys based on contextKey
return {
[contextKey]: useStoreHook,
[`${contextKey}Store`]: useStoreStoreHook,
} as {
[P in K]: typeof useStoreHook;
} & {
[P in `${K}Store`]: typeof useStoreStoreHook;
};
}