refactor(utils): extract redux extension logic into wrapper functions

src/utils/redux-extension/createReduxConnection.ts

@@ -0,0 +1,46 @@
+import { Message } from '../types';
+import { ReduxExtension } from './getReduxExtension';
+
+// Original but incomplete type of the redux extension package
+type ConnectResponse = ReturnType<NonNullable<ReduxExtension>['connect']>;
+
+export type Connection = {
+  /** Mark the connection as not initiated, so it can be initiated before using it. */
+  shouldInit?: boolean;
+
+  /** Initiate the connection and add it to the extension connections.
+   *  Should only be executed once in the live time of the connection.
+   */
+  init: ConnectResponse['init'];
+
+  // FIXME https://github.com/reduxjs/redux-devtools/issues/1097
+  /** Add a subscription to the connection.
+   *  The provided listener will be executed when the user interacts with the extension
+   *  with actions like time traveling, importing a state or the likes.
+   *
+   *  @param listener function to be executed when an action is submitted
+   *  @returns function to unsubscribe the applied listener
+   */
+  subscribe: (listener: (message: Message) => void) => (() => void) | undefined;
+
+  /** Send a new action to the connection to display the state change in the extension.
+   *  For example when the value of the store changes.
+   */
+  send: ConnectResponse['send'];
+};
+
+/** Wrapper for creating connections to the redux extension
+ *  Connections are used to display the stores value and value changes within the extension
+ *  as well as reacting to extension actions like time traveling.
+ **/
+export const createReduxConnection = (
+  extension: ReduxExtension | undefined,
+  name: string,
+) => {
+  if (!extension) return undefined;
+  const connection = extension.connect({ name });
+
+  return Object.assign(connection, {
+    shouldInit: true,
+  }) as Connection;
+};

src/utils/redux-extension/getReduxExtension.ts

@@ -0,0 +1,40 @@
+// Original but incomplete type of the redux extension package
+type Extension = NonNullable<typeof window.__REDUX_DEVTOOLS_EXTENSION__>;
+
+export type ReduxExtension = {
+  /** Create a connection to the extension.
+   *  This will connect a store (like an atom) to the extension and
+   *  display it within the extension tab.
+   *
+   *  @param options https://github.com/reduxjs/redux-devtools/blob/main/extension/docs/API/Arguments.md
+   *  @returns https://github.com/reduxjs/redux-devtools/blob/main/extension/docs/API/Methods.md#connectoptions
+   */
+  connect: Extension['connect'];
+
+  /** Disconnects all existing connections to the redux extension.
+   *  Only use this when you are sure that no other connection exists
+   *  or you want to remove all existing connections.
+   */
+  disconnect?: () => void;
+
+  /** Have a look at the documentation for more methods:
+   *  https://github.com/reduxjs/redux-devtools/blob/main/extension/docs/API/Methods.md
+   */
+};
+
+/** Returns the global redux extension object if available */
+export const getReduxExtension = (
+  enabled = __DEV__,
+): ReduxExtension | undefined => {
+  if (!enabled) {
+    return undefined;
+  }
+
+  const reduxExtension = window.__REDUX_DEVTOOLS_EXTENSION__;
+  if (!reduxExtension && __DEV__) {
+    console.warn('Please install/enable Redux devtools extension');
+    return undefined;
+  }
+
+  return reduxExtension;
+};

src/utils/useAtomDevtools.ts

@@ -1,7 +1,11 @@
 import { useEffect, useRef } from 'react';
 import { useAtom } from 'jotai/react';
 import type { Atom, WritableAtom } from 'jotai/vanilla';
-import { Message } from './types';
+import {
+  Connection,
+  createReduxConnection,
+} from './redux-extension/createReduxConnection';
+import { getReduxExtension } from './redux-extension/getReduxExtension';
 
 type DevtoolOptions = Parameters<typeof useAtom>[1] & {
   name?: string;
@@ -14,31 +18,13 @@ export function useAtomDevtools<Value, Result>(
 ): void {
   const { enabled, name } = options || {};
 
-  let extension: typeof window['__REDUX_DEVTOOLS_EXTENSION__'] | false;
-
-  try {
-    extension = (enabled ?? __DEV__) && window.__REDUX_DEVTOOLS_EXTENSION__;
-  } catch {
-    // ignored
-  }
-
-  if (!extension) {
-    if (__DEV__ && enabled) {
-      console.warn('Please install/enable Redux devtools extension');
-    }
-  }
+  const extension = getReduxExtension(enabled);
 
   const [value, setValue] = useAtom(anAtom, options);
 
   const lastValue = useRef(value);
   const isTimeTraveling = useRef(false);
-  const devtools = useRef<
-    ReturnType<
-      NonNullable<typeof window['__REDUX_DEVTOOLS_EXTENSION__']>['connect']
-    > & {
-      shouldInit?: boolean;
-    }
-  >();
+  const devtools = useRef<Connection>();
 
   const atomName = name || anAtom.debugLabel || anAtom.toString();
 
@@ -57,16 +43,9 @@ export function useAtomDevtools<Value, Result>(
       );
     };
 
-    devtools.current = extension.connect({ name: atomName });
+    devtools.current = createReduxConnection(extension, atomName);
 
-    const unsubscribe = (
-      devtools.current as unknown as {
-        // FIXME https://github.com/reduxjs/redux-devtools/issues/1097
-        subscribe: (
-          listener: (message: Message) => void,
-        ) => (() => void) | undefined;
-      }
-    ).subscribe((message) => {
+    const unsubscribe = devtools.current?.subscribe((message) => {
       if (message.type === 'ACTION' && message.payload) {
         try {
           setValueIfWritable(JSON.parse(message.payload));
@@ -106,7 +85,7 @@ export function useAtomDevtools<Value, Result>(
         });
       }
     });
-    devtools.current.shouldInit = true;
+
     return unsubscribe;
   }, [anAtom, extension, atomName, setValue]);
 

src/utils/useAtomsDevtools.ts

@@ -1,6 +1,10 @@
 import { useEffect, useRef } from 'react';
 import { AnyAtom, AnyAtomValue, AtomsSnapshot, Options } from '../types';
-import { Message } from './types';
+import {
+  Connection,
+  createReduxConnection,
+} from './redux-extension/createReduxConnection';
+import { getReduxExtension } from './redux-extension/getReduxExtension';
 import { useAtomsSnapshot } from './useAtomsSnapshot';
 import { useGotoAtomsSnapshot } from './useGotoAtomsSnapshot';
 
@@ -32,33 +36,15 @@ export function useAtomsDevtools(
 ): void {
   const { enabled } = options || {};
 
-  let extension: typeof window['__REDUX_DEVTOOLS_EXTENSION__'] | false;
-
-  try {
-    extension = (enabled ?? __DEV__) && window.__REDUX_DEVTOOLS_EXTENSION__;
-  } catch {
-    // ignored
-  }
-
-  if (!extension) {
-    if (__DEV__ && enabled) {
-      console.warn('Please install/enable Redux devtools extension');
-    }
-  }
+  const extension = getReduxExtension(enabled);
 
   // This an exception, we don't usually use utils in themselves!
   const atomsSnapshot = useAtomsSnapshot(options);
   const goToSnapshot = useGotoAtomsSnapshot(options);
 
   const isTimeTraveling = useRef(false);
   const isRecording = useRef(true);
-  const devtools = useRef<
-    ReturnType<
-      NonNullable<typeof window['__REDUX_DEVTOOLS_EXTENSION__']>['connect']
-    > & {
-      shouldInit?: boolean;
-    }
-  >();
+  const devtools = useRef<Connection>();
 
   const snapshots = useRef<AtomsSnapshot[]>([]);
 
@@ -74,16 +60,10 @@ export function useAtomsDevtools(
       }
       return snapshot;
     };
-    const connection = extension.connect({ name });
-
-    const devtoolsUnsubscribe = (
-      connection as unknown as {
-        // FIXME https://github.com/reduxjs/redux-devtools/issues/1097
-        subscribe: (
-          listener: (message: Message) => void,
-        ) => (() => void) | undefined;
-      }
-    ).subscribe((message) => {
+
+    devtools.current = createReduxConnection(extension, name);
+
+    const devtoolsUnsubscribe = devtools.current?.subscribe((message) => {
       switch (message.type) {
         case 'DISPATCH':
           switch (message.payload?.type) {
@@ -92,7 +72,7 @@ export function useAtomsDevtools(
               break;
 
             case 'COMMIT':
-              connection.init(getDevtoolsState(getSnapshotAt()));
+              devtools.current?.init(getDevtoolsState(getSnapshotAt()));
               snapshots.current = [];
               break;
 
@@ -109,10 +89,8 @@ export function useAtomsDevtools(
       }
     });
 
-    devtools.current = connection;
-    devtools.current.shouldInit = true;
     return () => {
-      (extension as any).disconnect();
+      extension?.disconnect?.();
       devtoolsUnsubscribe?.();
     };
   }, [extension, goToSnapshot, name]);