feat: support store v2 + drop support for older store methods

jotaijs · Apr 14, 2024 · 6bf0996 · 6bf0996
1 parent e8c5e96
commit 6bf0996
Show file tree

Hide file tree

Showing 11 changed files with 326 additions and 94 deletions.
diff --git a/src/DevTools/Extension/Extension.tsx b/src/DevTools/Extension/Extension.tsx
@@ -4,7 +4,7 @@ import clsx from 'clsx';
 import { useAtom, useSetAtom } from 'jotai/react';
 import { Store } from '../../types';
 import { isShellOpenAtom } from '../atoms/is-shell-open-atom';
-import { useSetCustomStore } from '../atoms/user-custom-store';
+import { useSetUserStore } from '../atoms/user-custom-store';
 import { useThemeMode } from '../hooks/useThemeMode';
 import { useDevtoolsJotaiStoreOptions } from '../internal-jotai-store';
 import { logo } from './assets/logo';
@@ -77,7 +77,7 @@ export const Extension = ({
     useDevtoolsJotaiStoreOptions(),
   );
 
-  const setUserStore = useSetCustomStore();
+  const setUserStore = useSetUserStore();
 
   React.useEffect(() => {
     setUserStore(store);

diff --git a/src/DevTools/Extension/components/Shell/components/AtomViewer/hooks/useInternalAtomValue.ts b/src/DevTools/Extension/components/Shell/components/AtomViewer/hooks/useInternalAtomValue.ts
@@ -7,7 +7,10 @@ import {
   useInternalStore,
 } from '../../../../../../../DevTools/internal-jotai-store';
 import { selectedAtomAtom } from '../atoms';
-import { useUserStore } from './../../../../../../hooks/useUserStore';
+import {
+  isDevToolsStore,
+  useUserStore,
+} from './../../../../../../hooks/useUserStore';
 import { isPromise, use } from './use';
 
 type Store = ReturnType<typeof useStore>;
@@ -63,11 +66,7 @@ export function useInternalAtomValue<Value>(atom: Atom<Value>) {
   }
 
   useEffect(() => {
-    const devSubscribeStore: Store['dev_subscribe_store'] =
-      // @ts-expect-error dev_subscribe_state is deprecated in <= 2.0.3
-      userStore?.dev_subscribe_store || userStore?.dev_subscribe_state;
-
-    if (!devSubscribeStore) return;
+    if (!isDevToolsStore(userStore)) return;
 
     const atomSubCb = <AtomSubscribeFunction>(() => {
       if (typeof delay === 'number') {
@@ -81,18 +80,18 @@ export function useInternalAtomValue<Value>(atom: Atom<Value>) {
     atomSubCb.isJotaiDevTools = true;
 
     const devSubCb = (
-      type?: Parameters<Parameters<typeof devSubscribeStore>[0]>[0],
+      action: Parameters<
+        Parameters<(typeof userStore)['subscribeStore']>[0]
+      >[0],
     ) => {
-      const normalizedType = typeof type === 'string' ? type : type?.type;
-
-      if (normalizedType !== 'unsub') {
+      if (action.type !== 'unsub') {
         return;
       }
 
       const activeValue = internalStore.get(selectedAtomAtom);
       if (activeValue) {
         const { l = [], t } =
-          userStore.dev_get_mounted?.(activeValue.atom) || {};
+          userStore.getMountedAtomState(activeValue.atom) || {};
         const listenersArray = Array.from(l);
         const areAllCallbacksInternal = listenersArray.every(
           isInternalAtomSubscribeFunction,
@@ -104,7 +103,8 @@ export function useInternalAtomValue<Value>(atom: Atom<Value>) {
         }
       }
     };
-    const devUnsubscribeStore = devSubscribeStore?.(devSubCb, 2);
+
+    const devUnsubscribeStore = userStore.subscribeStore(devSubCb);
 
     const unsub = userStore.sub(atom, atomSubCb);
     rerender();

diff --git a/src/DevTools/Extension/components/Shell/components/TimeTravel/useSyncSnapshotHistory.ts b/src/DevTools/Extension/components/Shell/components/TimeTravel/useSyncSnapshotHistory.ts
@@ -5,7 +5,10 @@ import {
   AtomsValues,
 } from '../../../../../../types';
 import { useDevToolsOptionsValue } from '../../../../../atoms/devtools-options';
-import { useUserStore } from '../../../../../hooks/useUserStore';
+import {
+  isDevToolsStore,
+  useUserStore,
+} from '../../../../../hooks/useUserStore';
 import { useInternalStore } from '../../../../../internal-jotai-store';
 import { atomToPrintable } from '../../../../../utils';
 import { createTimestamp } from '../../../../../utils/create-timestamp';
@@ -39,7 +42,8 @@ export default function useSyncSnapshotHistory() {
     useDevToolsOptionsValue();
 
   useEffect(() => {
-    if (!userStore?.dev_subscribe_store) return;
+    if (!isDevToolsStore(userStore)) return;
+
     const addToHistoryStack = (
       nextSnapshot: AtomsSnapshot,
       displayValues: SnapshotHistory['displayValues'],
@@ -73,8 +77,8 @@ export default function useSyncSnapshotHistory() {
       const values: AtomsValues = new Map();
       const displayValues: SnapshotHistory['displayValues'] = {};
       const dependents: AtomsDependents = new Map();
-      for (const atom of userStore.dev_get_mounted_atoms?.() || []) {
-        const atomState = userStore.dev_get_atom_state?.(atom);
+      for (const atom of userStore.getMountedAtoms() || []) {
+        const atomState = userStore.getAtomState(atom);
         if (atomState) {
           if ('v' in atomState) {
             values.set(atom, atomState.v);
@@ -106,32 +110,27 @@ export default function useSyncSnapshotHistory() {
             }
           }
         }
-        const mounted = userStore.dev_get_mounted?.(atom);
+        const mounted = userStore.getMountedAtomState(atom);
         if (mounted) {
           dependents.set(atom, mounted.t);
         }
       }
       addToHistoryStack({ values, dependents }, displayValues);
     };
     const cb = (
-      type?: Parameters<Parameters<typeof userStore.dev_subscribe_store>[0]>[0],
+      action: Parameters<Parameters<typeof userStore.subscribeStore>[0]>[0],
     ) => {
-      const normalizedType = typeof type === 'string' ? type : type?.type;
-      const isState =
-        normalizedType === 'write' ||
-        normalizedType === 'async-write' ||
-        // @ts-expect-error 'state' type is deprecated but we still support it for now
-        normalizedType === 'state';
+      const isWrite = action.type === 'set';
       if (
-        isState &&
+        isWrite &&
         shouldRecordSnapshotHistory &&
         !store.get(isTimeTravelingAtom)
       ) {
         collectValues();
       }
     };
 
-    const unsub = userStore.dev_subscribe_store?.(cb, 2);
+    const unsub = userStore.subscribeStore(cb);
     return unsub;
   }, [
     store,

diff --git a/src/DevTools/atoms/user-custom-store.ts b/src/DevTools/atoms/user-custom-store.ts
@@ -3,10 +3,10 @@ import { atom } from 'jotai/vanilla';
 import { Store } from 'src/types';
 import { useDevtoolsJotaiStoreOptions } from '../internal-jotai-store';
 
-const userCustomStoreAtom = atom<Store | undefined>(undefined);
+const userStore = atom<Store | undefined>(undefined);
 
-export const useUserCustomStoreValue = () =>
-  useAtomValue(userCustomStoreAtom, useDevtoolsJotaiStoreOptions());
+export const useUserStoreValue = () =>
+  useAtomValue(userStore, useDevtoolsJotaiStoreOptions());
 
-export const useSetCustomStore = () =>
-  useSetAtom(userCustomStoreAtom, useDevtoolsJotaiStoreOptions());
+export const useSetUserStore = () =>
+  useSetAtom(userStore, useDevtoolsJotaiStoreOptions());
diff --git a/src/DevTools/hooks/useUserStore.ts b/src/DevTools/hooks/useUserStore.ts
@@ -1,18 +1,21 @@
-import { useStore } from 'jotai/react';
 import { Options } from 'src/types';
-import { useUserCustomStoreValue } from '../atoms/user-custom-store';
+import { useDevToolsStore } from '../../utils/hooks/useDevToolsStore';
+import { useUserStoreValue } from '../atoms/user-custom-store';
+import { isDevToolsStore } from './../../utils/internals/compose-with-devtools';
 
-export const useUserStore = () => {
-  const possibleUserStore = useUserCustomStoreValue();
+export const useUserStore = (): ReturnType<typeof useDevToolsStore> => {
+  const possibleUserStore = useUserStoreValue();
 
-  const userStore = useStore(
+  const userStore = useDevToolsStore(
     // This defaults to user's default store in a `provider-less` mode
     possibleUserStore ? { store: possibleUserStore } : undefined,
   );
 
   return userStore;
 };
 
+export { isDevToolsStore };
+
 export const useUserStoreOptions = (): Options => {
   const store = useUserStore();
   return { store: store };

diff --git a/src/types.ts b/src/types.ts
@@ -1,12 +1,26 @@
 import { useStore } from 'jotai/react';
-import type { Atom, createStore } from 'jotai/vanilla';
+import type {
+  Atom,
+  WritableAtom,
+  createStore as createStoreV1,
+} from 'jotai/vanilla';
+import type {
+  INTERNAL_DevStoreRev4,
+  INTERNAL_PrdStore,
+} from 'jotai/vanilla/store2';
 
-export type Store = ReturnType<typeof createStore>;
+export type StoreV1 = ReturnType<typeof createStoreV1>;
+export type StoreV2 = INTERNAL_DevStoreRev4 & INTERNAL_PrdStore;
+
+export type Store = StoreV1 | StoreV2;
 
 export type Options = Parameters<typeof useStore>[0];
 
 export type AnyAtomValue = unknown;
+export type AnyAtomError = unknown;
 export type AnyAtom = Atom<AnyAtomValue>;
+export type AnyWritableAtom = WritableAtom<AnyAtomValue, unknown[], unknown>;
+
 export type AtomsValues = Map<AnyAtom, AnyAtomValue>; // immutable
 export type AtomsDependents = Map<AnyAtom, Set<AnyAtom>>; // immutable
 export type AtomsSnapshot = Readonly<{

diff --git a/src/utils/hooks/useDevToolsStore.ts b/src/utils/hooks/useDevToolsStore.ts
@@ -0,0 +1,16 @@
+import { useStore } from 'jotai';
+import { Options } from '../../types';
+import {
+  composeWithDevTools,
+  isDevToolsStore,
+} from '../internals/compose-with-devtools';
+
+export const useDevToolsStore = (
+  options: Options,
+): ReturnType<typeof composeWithDevTools> => {
+  const store = useStore(options);
+
+  return composeWithDevTools(store);
+};
+
+export { isDevToolsStore };