diff --git a/src/createIsolation.tsx b/src/createIsolation.tsx index a4c677c..784905f 100644 --- a/src/createIsolation.tsx +++ b/src/createIsolation.tsx @@ -6,6 +6,7 @@ import { useAtom as useAtomOrig, useAtomValue as useAtomValueOrig, useSetAtom as useSetAtomOrig, + useStore as useStoreOrig, } from 'jotai/react'; import { useHydrateAtoms } from 'jotai/react/utils'; @@ -36,29 +37,26 @@ export function createIsolation() { ); }; - const useAtom = ((anAtom: any, options?: any) => { + const useStore = ((options?: any) => { const store = useContext(StoreContext); - if (!store) { - throw new Error('Missing Provider from createIsolation'); - } + if (!store) throw new Error('Missing Provider from createIsolation'); + return useStoreOrig({ store, ...options }); + }) as typeof useStoreOrig; + + const useAtom = ((anAtom: any, options?: any) => { + const store = useStore(); return useAtomOrig(anAtom, { store, ...options }); }) as typeof useAtomOrig; const useAtomValue = ((anAtom: any, options?: any) => { - const store = useContext(StoreContext); - if (!store) { - throw new Error('Missing Provider from createIsolation'); - } + const store = useStore(); return useAtomValueOrig(anAtom, { store, ...options }); }) as typeof useAtomValueOrig; const useSetAtom = ((anAtom: any, options?: any) => { - const store = useContext(StoreContext); - if (!store) { - throw new Error('Missing Provider from createIsolation'); - } + const store = useStore(); return useSetAtomOrig(anAtom, { store, ...options }); }) as typeof useSetAtomOrig; - return { Provider, useAtom, useAtomValue, useSetAtom }; + return { Provider, useStore, useAtom, useAtomValue, useSetAtom }; }