Refactor atomWithFn into util file

packages/jotai-x/src/atomWithFn.ts

@@ -0,0 +1,32 @@
+import { atom } from 'jotai';
+
+import type { WritableAtom } from 'jotai/vanilla';
+
+type WrapFn<T> = T extends (...args: infer _A) => infer _R ? { __fn: T } : T;
+
+const wrapFn = <T>(fnOrValue: T): WrapFn<T> =>
+  (typeof fnOrValue === 'function' ? { __fn: fnOrValue } : fnOrValue) as any;
+
+type UnwrapFn<T> = T extends { __fn: infer U } ? U : T;
+
+const unwrapFn = <T>(wrappedFnOrValue: T): UnwrapFn<T> =>
+  (wrappedFnOrValue &&
+  typeof wrappedFnOrValue === 'object' &&
+  '__fn' in wrappedFnOrValue
+    ? wrappedFnOrValue.__fn
+    : wrappedFnOrValue) as any;
+
+/**
+ * Jotai atoms don't allow functions as values by default. This function is a
+ * drop-in replacement for `atom` that wraps functions in an object while
+ * leaving non-functions unchanged. The wrapper object should be completely
+ * invisible to consumers of the atom.
+ */
+export const atomWithFn = <T>(initialValue: T): WritableAtom<T, [T], void> => {
+  const baseAtom = atom(wrapFn(initialValue));
+
+  return atom(
+    (get) => unwrapFn(get(baseAtom)) as T,
+    (_get, set, value) => set(baseAtom, wrapFn(value))
+  );
+};

packages/jotai-x/src/createAtomStore.ts

@@ -1,6 +1,7 @@
-import { atom, useAtom, useAtomValue, useSetAtom } from 'jotai';
+import { useAtom, useAtomValue, useSetAtom } from 'jotai';
 import { useHydrateAtoms } from 'jotai/utils';
 
+import { atomWithFn } from './atomWithFn';
 import { createAtomProvider, useAtomStore } from './createAtomProvider';
 
 import type { ProviderProps } from './createAtomProvider';
@@ -17,30 +18,6 @@ export type UseAtomOptions = {
 
 type UseAtomOptionsOrScope = UseAtomOptions | string;
 
-// Jotai does not support functions in atoms, so wrap functions in objects
-type WrapFn<T> = T extends (...args: infer _A) => infer _R ? { __fn: T } : T;
-
-const wrapFn = <T>(fnOrValue: T): WrapFn<T> =>
-  (typeof fnOrValue === 'function' ? { __fn: fnOrValue } : fnOrValue) as any;
-
-type UnwrapFn<T> = T extends { __fn: infer U } ? U : T;
-
-const unwrapFn = <T>(wrappedFnOrValue: T): UnwrapFn<T> =>
-  (wrappedFnOrValue &&
-  typeof wrappedFnOrValue === 'object' &&
-  '__fn' in wrappedFnOrValue
-    ? wrappedFnOrValue.__fn
-    : wrappedFnOrValue) as any;
-
-const atomWithFn = <T>(initialValue: T): SimpleWritableAtom<T> => {
-  const baseAtom = atom(wrapFn(initialValue));
-
-  return atom(
-    (get) => unwrapFn(get(baseAtom)) as T,
-    (_get, set, value) => set(baseAtom, wrapFn(value))
-  );
-};
-
 type GetRecord<O> = {
   [K in keyof O]: O[K] extends Atom<infer V>
     ? (options?: UseAtomOptionsOrScope) => V

packages/jotai-x/src/index.ts

@@ -3,6 +3,7 @@
  */
 
 export * from './atomProvider';
+export * from './atomWithFn';
 export * from './createAtomProvider';
 export * from './createAtomStore';
 export * from './useHydrateStore';