feat: withAtomEffect allows binding effects to atoms

README.md

@@ -12,7 +12,7 @@ npm i jotai-effect
 
 `atomEffect` is a utility function for declaring side effects and synchronizing atoms in Jotai. It is useful for observing and reacting to state changes.
 
-## Parameters
+### Parameters
 
 ```ts
 type CleanupFn = () => void
@@ -27,7 +27,7 @@ declare function atomEffect(effectFn: EffectFn): Atom<void>
 
 **effectFn** (required): A function for listening to state updates with `get` and writing state updates with `set`. The `effectFn` is useful for creating side effects that interact with other Jotai atoms. You can cleanup these side effects by returning a cleanup function.
 
-## Usage
+### Usage
 
 Subscribe to Atom Changes
 
@@ -54,7 +54,7 @@ const subscriptionEffect = atomEffect((get, set) => {
 })
 ```
 
-## Mounting with Atoms or Hooks
+### Mounting with Atoms or Hooks
 
 After defining an effect using `atomEffect`, it can be integrated within another atom's read function or passed to Jotai hooks.
 
@@ -74,7 +74,7 @@ function MyComponent() {
 
 <CodeSandbox id="tg9xsf" />
 
-## The `atomEffect` behavior
+### The `atomEffect` behavior
 
 - **Cleanup Function:**
   The cleanup function is invoked on unmount or before re-evaluation.
@@ -371,3 +371,35 @@ atomEffects are distinguished from useEffect in a few other ways. They can direc
 
 Both useEffect and atomEffect have their own advantages and applications. Your project’s specific needs and your comfort level should guide your selection.
 Always lean towards an approach that gives you a smoother, more intuitive development experience. Happy coding!
+
+## withAtomEffect
+
+`withAtomEffect` is a utility to define an effect that is bound to the target atom. This is useful for creating effects that are active when the target atom is mounted.
+
+### Parameters
+
+```ts
+declare function withAtomEffect<T>(
+  targetAtom: Atom<T>,
+  effectFn: EffectFn,
+): Atom<T>
+```
+
+**targetAtom** (required): The atom to which the effect is mounted.
+
+**effectFn** (required): A function for listening to state updates with `get` and writing state updates with `set`.
+
+**Returns:** An atom that is equivalent to the target atom but with the effect mounted.
+
+### Usage
+
+```js
+import { withAtomEffect } from 'jotai-effect'
+
+const anAtom = atom(0)
+const loggingAtom = withAtomEffect(anAtom, (get, set) => {
+  // runs on mount or whenever anAtom changes
+  const value = get(anAtom)
+  loggingService.setValue(value)
+})
+```

__tests__/atomEffect.strict.test.ts

@@ -3,6 +3,7 @@ import { act, renderHook, waitFor } from '@testing-library/react'
 import { useAtom, useAtomValue, useSetAtom } from 'jotai/react'
 import { atom, getDefaultStore } from 'jotai/vanilla'
 import { atomEffect } from '../src/atomEffect'
+import { assert, delay, increment, incrementLetter } from './test-utils'
 
 const wrapper = StrictMode
 
@@ -623,23 +624,3 @@ it('should not run the effect when the effectAtom is unmounted', async () => {
   await act(() => setCount(increment))
   expect(runCount).toBe(2)
 })
-
-function increment(count: number) {
-  return count + 1
-}
-
-function incrementLetter(str: string) {
-  return String.fromCharCode(increment(str.charCodeAt(0)))
-}
-
-function delay(ms: number) {
-  return new Promise((resolve) => {
-    setTimeout(resolve, ms)
-  })
-}
-
-function assert(value: boolean, message?: string): asserts value {
-  if (!value) {
-    throw new Error(message ?? 'assertion failed')
-  }
-}

__tests__/atomEffect.test.tsx

@@ -1,13 +1,15 @@
-import React, {
-  Component,
-  type ErrorInfo,
-  type ReactNode,
-  useEffect,
-} from 'react'
+import React, { useEffect } from 'react'
 import { act, render, renderHook, waitFor } from '@testing-library/react'
 import { Provider, useAtom, useAtomValue, useSetAtom } from 'jotai/react'
 import { type Atom, atom, createStore, getDefaultStore } from 'jotai/vanilla'
 import { atomEffect } from '../src/atomEffect'
+import {
+  ErrorBoundary,
+  assert,
+  delay,
+  increment,
+  incrementLetter,
+} from './test-utils'
 
 type AnyAtom = Atom<any>
 
@@ -948,7 +950,7 @@ it('should trigger an error boundary when an error is thrown in a cleanup', asyn
     return (
       <Provider store={store}>
         <ErrorBoundary
-          componentDidCatch={(error: Error, _errorInfo: ErrorInfo) => {
+          componentDidCatch={(error, _errorInfo) => {
             if (!didThrow) {
               expect(error.message).toBe('effect cleanup error')
             }
@@ -964,49 +966,3 @@ it('should trigger an error boundary when an error is thrown in a cleanup', asyn
   act(() => store.set(refreshAtom, increment))
   await waitFor(() => assert(didThrow))
 })
-
-function increment(count: number) {
-  return count + 1
-}
-
-function incrementLetter(str: string) {
-  return String.fromCharCode(increment(str.charCodeAt(0)))
-}
-
-function delay(ms: number) {
-  return new Promise((resolve) => {
-    setTimeout(resolve, ms)
-  })
-}
-
-function assert(value: boolean, message?: string): asserts value {
-  if (!value) {
-    throw new Error(message ?? 'assertion failed')
-  }
-}
-
-type ErrorBoundaryState = {
-  hasError: boolean
-}
-type ErrorBoundaryProps = {
-  componentDidCatch?: (error: Error, errorInfo: ErrorInfo) => void
-  children: ReactNode
-}
-class ErrorBoundary extends Component<ErrorBoundaryProps, ErrorBoundaryState> {
-  state = { hasError: false }
-
-  static getDerivedStateFromError(): ErrorBoundaryState {
-    return { hasError: true }
-  }
-
-  componentDidCatch(error: Error, _errorInfo: ErrorInfo): void {
-    this.props.componentDidCatch?.(error, _errorInfo)
-  }
-
-  render() {
-    if (this.state.hasError) {
-      return <div>error</div>
-    }
-    return this.props.children
-  }
-}

__tests__/test-utils.ts

@@ -0,0 +1,50 @@
+import { Component, type ErrorInfo, type ReactNode, createElement } from 'react'
+
+export function increment(count: number) {
+  return count + 1
+}
+
+export function incrementLetter(str: string) {
+  return String.fromCharCode(increment(str.charCodeAt(0)))
+}
+
+export function delay(ms: number) {
+  return new Promise((resolve) => {
+    setTimeout(resolve, ms)
+  })
+}
+
+export function assert(value: boolean, message?: string): asserts value {
+  if (!value) {
+    throw new Error(message ?? 'assertion failed')
+  }
+}
+
+type ErrorBoundaryState = {
+  hasError: boolean
+}
+type ErrorBoundaryProps = {
+  componentDidCatch?: (error: Error, errorInfo: ErrorInfo) => void
+  children: ReactNode
+}
+export class ErrorBoundary extends Component<
+  ErrorBoundaryProps,
+  ErrorBoundaryState
+> {
+  state = { hasError: false }
+
+  static getDerivedStateFromError(): ErrorBoundaryState {
+    return { hasError: true }
+  }
+
+  componentDidCatch(error: Error, _errorInfo: ErrorInfo): void {
+    this.props.componentDidCatch?.(error, _errorInfo)
+  }
+
+  render() {
+    if (this.state.hasError) {
+      return createElement('div', { children: 'error' })
+    }
+    return this.props.children
+  }
+}

__tests__/withAtomEffect.test.ts

@@ -0,0 +1,61 @@
+import { atom, createStore } from 'jotai/vanilla'
+import { withAtomEffect } from '../src/withAtomEffect'
+import { delay } from './test-utils'
+
+describe('withAtomEffect', () => {
+  it('ensures readonly atoms remain readonly', () => {
+    const readOnlyAtom = atom(() => 10)
+    const enhancedAtom = withAtomEffect(readOnlyAtom, () => {})
+    const store = createStore()
+    store.sub(enhancedAtom, () => {})
+    expect(store.get(enhancedAtom)).toBe(10)
+    expect(() => {
+      // @ts-expect-error: should error
+      store.set(enhancedAtom, 20)
+    }).toThrow()
+  })
+
+  it('ensures writable atoms remain writable', () => {
+    const writableAtom = atom(0)
+    const enhancedAtom = withAtomEffect(writableAtom, () => {})
+    const store = createStore()
+    store.sub(enhancedAtom, () => {})
+    store.set(enhancedAtom, 5)
+    expect(store.get(enhancedAtom)).toBe(5)
+  })
+
+  it('calls effect on initial use and on dependencies change', async () => {
+    const baseAtom = atom(0)
+    const effectMock = jest.fn()
+    const enhancedAtom = withAtomEffect(baseAtom, (get) => {
+      effectMock()
+      get(baseAtom)
+    })
+    const store = createStore()
+    store.sub(enhancedAtom, () => {})
+    await delay(0)
+    expect(effectMock).toHaveBeenCalledTimes(1)
+    store.set(baseAtom, 1)
+    await delay(0)
+    expect(effectMock).toHaveBeenCalledTimes(2)
+  })
+
+  it('cleans up when the atom is no longer in use', async () => {
+    const cleanupMock = jest.fn()
+    const baseAtom = atom(0)
+    const mountMock = jest.fn()
+    const enhancedAtom = withAtomEffect(baseAtom, () => {
+      mountMock()
+      return () => {
+        cleanupMock()
+      }
+    })
+    const store = createStore()
+    const unsubscribe = store.sub(enhancedAtom, () => {})
+    await delay(0)
+    expect(mountMock).toHaveBeenCalledTimes(1)
+    unsubscribe()
+    await delay(0)
+    expect(cleanupMock).toHaveBeenCalledTimes(1)
+  })
+})

package.json

@@ -36,7 +36,10 @@
   },
   "jest": {
     "testEnvironment": "jsdom",
-    "preset": "ts-jest/presets/js-with-ts"
+    "preset": "ts-jest/presets/js-with-ts",
+    "testMatch": [
+      "**/__tests__/**/*.test.ts?(x)"
+    ]
   },
   "keywords": [
     "jotai",
@@ -64,7 +67,7 @@
     "html-webpack-plugin": "^5.5.3",
     "jest": "^29.7.0",
     "jest-environment-jsdom": "^29.7.0",
-    "jotai": "^2.9.3",
+    "jotai": "2.6.3",
     "microbundle": "^0.15.1",
     "npm-run-all": "^4.1.5",
     "prettier": "^3.0.3",
@@ -79,6 +82,6 @@
     "webpack-dev-server": "^4.15.1"
   },
   "peerDependencies": {
-    "jotai": ">=2.5.0"
+    "jotai": ">=2.5.0 <2.9.0"
   }
 }

src/atomEffect.ts

@@ -4,6 +4,10 @@ import { atom } from 'jotai/vanilla'
 type CleanupFn = () => void
 type GetterWithPeek = Getter & { peek: Getter }
 type SetterWithRecurse = Setter & { recurse: Setter }
+export type EffectFn = (
+  get: GetterWithPeek,
+  set: SetterWithRecurse
+) => void | CleanupFn
 
 export function atomEffect(
   effectFn: (get: GetterWithPeek, set: SetterWithRecurse) => void | CleanupFn

src/withAtomEffect.ts

@@ -0,0 +1,28 @@
+import type { Atom, Getter, Setter, WritableAtom } from 'jotai/vanilla'
+import { atom } from 'jotai/vanilla'
+import type { EffectFn } from './atomEffect'
+import { atomEffect } from './atomEffect'
+
+export function withAtomEffect<Value, Args extends unknown[], Result>(
+  targetAtom: WritableAtom<Value, Args, Result>,
+  effectFn: EffectFn
+): WritableAtom<Value, Args, Result>
+
+export function withAtomEffect<Value>(
+  targetAtom: Atom<Value>,
+  effectFn: EffectFn
+): Atom<Value>
+
+export function withAtomEffect<Value, Args extends unknown[], Result>(
+  targetAtom: Atom<Value> | WritableAtom<Value, Args, Result>,
+  effectFn: EffectFn
+): Atom<Value> | WritableAtom<Value, Args, Result> {
+  const effect = atomEffect(effectFn)
+  const readFn = (get: Getter) => void get(effect) ?? get(targetAtom)
+  if ('write' in targetAtom) {
+    type WriteFn = (get: Getter, set: Setter, ...args: Args) => Result
+    const writeFn: WriteFn = (_, set, ...args) => set(targetAtom, ...args)
+    return atom(readFn, writeFn)
+  }
+  return atom(readFn)
+}