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.
@@ -342,9 +342,42 @@ Aside from mount events, the effect runs when any of its dependencies change val
 
   </details>
 
-### Comparison with useEffect
+## withAtomEffect
 
-#### Component Side Effects
+`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)
+})
+```
+
+
+## Comparison with useEffect
+
+### Component Side Effects
 
 [useEffect](https://react.dev/reference/react/useEffect) is a React Hook that lets you synchronize a component with an external system.
 
@@ -354,7 +387,7 @@ Each call to a hook has a completely isolated state.
 This isolation can be referred to as _component-scoped_.
 For synchronizing component props and state with a Jotai atom, you should use the useEffect hook.
 
-#### Global Side Effects
+### Global Side Effects
 
 For setting up global side-effects, deciding between useEffect and atomEffect comes down to developer preference.
 Whether you prefer to build this logic directly into the component or build this logic into the Jotai state model depends on what mental model you adopt.
@@ -367,7 +400,7 @@ The same guarantee can be achieved with the useEffect hook if you ensure that th
 
 atomEffects are distinguished from useEffect in a few other ways. They can directly react to atom state changes, are resistent to infinite loops, and can be mounted conditionally.
 
-#### It's up to you
+### It's up to you
 
 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!

__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