👻

React の状態管理ライブラリ Jotai の詳説と Tips

2024/04/27に公開

JavaScript

TypeScript

jotai

tech

概要

Jotai, primitive and flexible state management for React
Jotai は Recoil からインスピレーションを受けて開発されたグローバルステート管理ライブラリです。コア API が厳選されており良い意味で少なくて機能もシンプルなので Recoil を利用した事がある人であれば直ぐに理解できるでしょうし、そうでない人でも useState に近い使用感なので比較的直ぐに理解して使えるようになると思います。

他のグローバルステート管理ライブラリと共通する Context や useState だと起きやすい以下の問題を解決します。

props のバケツリレーによる依存関係や影響範囲の複雑化
多数の Provider を管理することによる複雑化
これらの影響による過剰なメモ化による再描画対策

Recoil と比較した特徴としては TypeScript で開発されているので型もドキュメント上の型に関する情報も充実しているのと現在も活発に開発されている点です。

理論周りを詳しく知りたい人向けのおすすめ記事を紹介しておきます。

Jotai のコンセプトや特徴をより詳しく知りたい
- Concepts — Jotai, primitive and flexible state management for React
- Comparison — Jotai, primitive and flexible state management for React
グローバルステート管理ライブラリが必要な背景を知りたい
- Facebook製の新しいステート管理ライブラリ「Recoil」を最速で理解する - uhyo/blog
- 前半部分が Recoil に限らない背景の説明になっています

本記事では Jotai のコア機能と Tips を筆者が公式ドキュメントではわかりづらかった点についてソースコードを読んだり、コードで試行錯誤して理解した内容を元に説明したいと思います。

インストール

pnpm add jotai

使い方の要約

各 API の詳細な説明に入る前に全体像を掴めていた方が理解しやすいと思うので、利用頻度の高い API を使って要点だけ説明しておきます。

jotai におけるグローバルステートは Atom という単位で管理するのでまずは Atom を定義します。

import { atom } from 'jotai'

// デフォルト値を指定して atom を作成します
const countAtom = atom(0)
// 型の指定もできます
const currentUserAtom = atom<User | undefined>(undefined)

Atom が定義できたらコンポーネントまたはカスタムフックで利用できます。

import { useAtom, useAtomValue, useSetAtom } from 'jotai'
import { countAtom } from './countAtom'

export const Example = () => {
  // 値と更新関数の取得
  const [count, setCount] = useAtom(countAtom)
  // 値のみ取得
  const count = useAtomValue(countAtom)
  // 更新関数のみ取得
  const setCount = useSetAtom(countAtom)
  ...
  // 更新関数の利用方法は useState と同じです
  setCount(10)
  setCount((prev) => prev + 1)
}

Jotai はデフォルトでプロバイダーレスモードになっており Provider なしで動きますが、 Provider を利用する事も可能です。Jotai の Provider は Context と同様に並列に並べたり入れ子にする事が可能で Provider 毎に値を保持します。

import { Provider } from 'jotai'

export const Example = () => {
  return (
    <>
      <Provider>
        <Counter /> // countAtom を利用したコンポーネント
        <Provider>
          <Counter /> // 上の Counter とは別管理のステートとなる
        </Provider>
      </Provider>
    </>
  )
}

Atom

atom — Jotai, primitive and flexible state management for React
atom 関数は atom config と呼ばれる Atom の定義オブジェクトを作成して返します。atom config はイミュータブルで Atom の値は保持していません。Atom の実際の値は後述する Store が保持しており、Store から Atom を参照する為に使うオブジェクトです。

似た名称があって対象がどれなのか混乱しやすいので改めて整理しますが、本ドキュメントでは以下の定義で書いています。

Atom
- 1 つのグローバルステートのこと
atom config
- Atom を参照する為の定義オブジェクトのこと
atom
- Atom を定義して atom config を作成する為の関数のこと

基本となる atom config の作成方法は前述の通りで、プリミティブな値を保持するのであれば引数にデフォルト値を指定するだけです。オブジェクトを保持させたい場合は型指定をする事ができます。

import { atom } from 'jotai'
const count = atom(0)
const userAtom = atom<User>({ id: 1, name: 'akineko' })

派生 Atom

atom 関数は既存の Atom を利用した以下の 3 種の派生 Atom を作成する事もできます。

Read-Write Atom
Read-only Atom
Write-only Atom

それぞれの派生 Atom の作成例は以下の通りです。

// 渡す関数の引数の get, set を使って値の取得や更新を行います
const readWriteAtom = atom(
  (get) => get(priceAtom) * 1.1,
  (get, set, newPrice) => set(priceAtom, newPrice / 1.1)
)

// Read-only を作る時は第二引数の関数を省略します
const readOnlyAtom = atom((get) => get(priceAtom) * get(countAtom))

// Write-only を作る際は第一引数に null を渡します
const writeOnlyAtom = atom(
  null,
  (get, set, count) => {
    set(countAtom, count)
    set(totalAtom, get(priceAtom) * get(coountAtom))
  }
)

コンポーネント/カスタムフック内での作成

atom config はコンポーネントやカスタムフック内でも作る事が出来ますが、その際には参照が保たれるように useMemo などでメモ化する必要があります。もしメモ化しなかった場合は無限ループが発生するので注意して下さい。

const Example = () => {
  const valueAtom = useMemo(() => atom(0), [])
}

振る舞い

最後に Atom の振る舞いについて説明します。Atom は初期状態では値は存在せず、初めて参照する時に定義されていた値を格納して返します。派生 Atom の場合も同様に初めて参照する時に派生元の Atom を参照し、計算した値を格納して返します。そして、Atom を参照しているコンポーネントが全て unmount されると値は GC されます。

Store / Provider

Store — Jotai, primitive and flexible state management for React
Provider — Jotai, primitive and flexible state management for React
Provider と Store は一緒に知った方が理解しやすいので合わせて説明します。

Store は Atom の値を格納するオブジェクトで、 Provider は内部に Store を保持して自身の子コンポーネントに提供します。Provider は内部で自動的に Store を作成しますが props 経由で Store を渡す事もできます。

Store は createStore 関数で作成することができ、下記は利用する場合の例です。

import { createStore, Provider } from 'jotai'

const store = createStore()

const Example = () => (
  <>
    // 自動生成される内部の Store を利用
    <Provider><Counter></Provider>
    // 外部から Store を提供
    <Provider store={store}><Counter></Provider>
  </>
)

Provider と Store の主な利用用途は以下の通りです。

Context のようにツリー毎に異なる状態を保持したい
Atom に動的な初期値を設定したい
コンポーネントの unmount により Atom をクリアしたい

Store オブジェクト

Store オブジェクトは以下のメソッドを提供しています。

get
- atom config を引数渡すとその Atom の値を取得できます
set
- 第一引数に atom config 第二引数に更新内容を指定して値を更新します
- 値の更新内容の指定は useState と同じ値もしくは関数です
sub
- いわゆる購読処理を登録する為の関数です
- 第一引数に atom config 第二引数に更新があった時に実行される関数を指定します
- 戻り値は呼び出すと購読処理を解除する関数です

import { atom, createStore } from 'jotai'

const countAtom = atom(0)

const store = createStore()
// 取得
store.get(countAtom) // 0
// 更新
store.set(countAtom, 1)
store.set(countAtom, (prev) => prev + 1)
// 購読
const unsub = store.sub(countAtom, () => console.log(store.get(countAtom)))
unsub()

デフォルト Store

プロバイダーレスモードの時に使用される jotai が内部で作成する Store は getDefaultStore 関数で取得できます。

import { getDefaultStore } from 'jotai'
const defaultStore = getDefaultStore()

useStore

直近の親コンポーネントが保持している Store は useStore でを取得できます。

import { useStore } from 'jotai'

const Example = () => {
  const store = useStore()
}

useAtom

useAtom — Jotai, primitive and flexible state management for React
jotai は Atom をコンポーネントで利用する為のコア API として useAtom を中心とした以下の 3 つのカスタムフックを提供しています。

useAtom
- 同様に対象となる Atom の値と更新関数を返します
- 戻り値の形式は useState と同様の [value, setValue] のタプルです
useAtomValue
- 対象となる Atom の値だけを返します
useSetAtom
- 対象となる Atom の更新関数だけを返します

これらのカスタムフックのインターフェイスは下記の通り共通しています。第一引数が atom config で第二引数が省略可能なオプションで執筆時点では Store のみです。

import { useAtom, useAtomValue, useSetAtom } from 'jotai'

const [value, setValue] = useAtom(countAtom, { store })
const value = useAtomValue(countAtom, { store })
const setValue = useSetAtom(countAtom, { store })

Store オプションを省略時の振る舞いは、直近の親となる Provider があればその Store が対象となり、なければデフォルトの Store が対象となります。つまり、複数の Provider が存在すると同じ Atom に対して Provider 毎に異なる値が保持されるようになります。アプリケーション全体で共通の値を利用したい Atom の場合は困るので、そういった Atom を使用する際はこの Store オプションを利用してコントロールします。

Tips

コア API の解説については以上です。
最後に応用的な使い方や個人的におすすめなユーティリティや拡張機能の使い方を紹介します。筆者が執筆時点で知っている情報なので、後日に追加で調べた事がありましたら追記していきます。

テスト方法

アプリケーション外での利用用途だとテスト用に初期値を設定したり、値の取得/更新が Store と Provider を利用した wrapper を用意するだけで簡単に行えます。私はよく以下のような wrapper 作成関数を用意しています。

type WrapperProps = { children: ReactNode }
export const createJotaiWrapper = () => {
  const store = createStore()
  const wrapper = ({ children }): WrapperProps) => (
    <Provider store={store}>{children}</Provider>
  )
  return { store, wrapper }
}

Testing — Jotai, primitive and flexible state management for React
個人的には Store を利用する方が柔軟に対応しやすいので好みですが、公式では別の方法で初期値を提供する方法も紹介されています。

Atom の動的な参照

Atoms in atom — Jotai
atom config は Atom や useState などにも持たせることが出来るので動的に参照する Atom を切り替えたりもできます。

雑ですが利用例です。

const pricesAtom = atom({
  pen: atom(10),
  pineapple: atom(20),
  apple: atom(15),
})

const Store = () => {
  const prices = useAtomValue(pricesAtom)
  // prices の中の特定の Atom を参照する State
  const [item, setItem] = useState(prices['pen'])
  // 上で参照している Atom の値を取得
  const price = useAtomValue(item)
  ...
  // prices を使った商品選択 UI や選択したアイテムの価格表示
}

実行時のみ Atom を参照する関数

Callback — Jotai, primitive and flexible state management for React
jotai/utils から提供されている useAtomCallback を利用すれば関数を呼び出した時にのみ Atom を参照して処理を行う関数を作成できるので、Atom の値が処理の実行時にしか必要ないコンポーネントの不要な再描画を防ぐ事ができます。

useAtomCallback は (get: Getter, set: Setter, ...args: Args) => Result の型の引数を受け取り、(...args: Args) => Result の型の関数を返すカスタムフックです。useAtomCallback に渡す関数は参照が保たれるようにグローバルに定義された関数もしくは useCallback を利用して作成した関数を渡して下さい。

import { useAtomCallback } from 'jotai/utils'

const Example = () => {
  const addItemToCart = useAtomCallback(
    useCallback((get, set, item, count) => {
      const price = get(pricesAtom)[item]
      set(totalPriceAtom, (prev) => prev + price * count)
      return get(totalPriceAtom)
    },[])
  )
  ...
  const currentTotal = addItemToCart('apple', 3)
}

Reducer

Reducer — Jotai, primitive and flexible state management for React
jotai/utils から提供されている atomWithReducer を利用すれば Atom レベルで Reducer パターンを適用できます。

export function atomWithReducer<Value, Action>(
  initialValue: Value,
  reducer: (value: Value, action: Action) => Value,
): WritableAtom<Value, [Action], void>

atomWithReducer の型定義は上記の通りで、型パラメータの第一引数が Value で Atom が保持する値の型、第二引数が Action でReducer 関数の第二引数の型です。

下記のコードは型パラメータを指定した例です。このコードだと型パラメータを省略して推論にまかせても問題ありませんが、Atom の値を User | undefined のようなユニオン型にしたい時には指定する必要があります。

import { atomWithReducer } from 'jotai/utils'

type Value = number
type Action = { type: 'Add' | 'Sub'; delta: number }

const reducer = (prevCount: Value, action: Action) => {
  switch (action.type) {
    case 'Add':
      return prevCount + action.delta
    case 'Sub':
      return prevCount - action.delta
    default:
      return prevCount
  }
}

const reducerAtom = atomWithReducer<Value, Action>(0, reducer)

atomWithReducer の atom config は useAtom などで使用すると set 関数の代わりに Action を引数に取る dispatch を返します。

const Example = () => {
  const [count, dispatch] = useAtom(reducerAtom)
  ...
  dispatch({ type: 'Add', delta: 5 })
}

useReducerAtom

useReducerAtom — Jotai
Atom レベルではなくカスタムフックレベルで Reducer パターンを適用したい場合のサンプルレシピとして公式より useReducerAtom が紹介されています。実は jotai/utils で提供されていたカスタムフックなのですが、現在は DEPRECATED となっており削除される予定です。

Is it posible to combine useSetAtom with useReducerAtom? · pmndrs/jotai · Discussion #2464
詳細はこちらの Discussion に書かれているのですが、この実装だと dispatch だけが必要なコンポーネントなども値が更新された時に再描画の対象となってしまうので、useSetAtom を使った実装の方がいいのと利用者も少ないのでレシピだけ公開して非推奨になったようです。

確かに Jotai を利用するという事は値の参照と更新を行うコンポーネントが離れているケースの方が多いので、レシピを参考に自身のニーズにあったものを作る方が良さそうです。

export function useReducerAtom<Value, Action>(
  anAtom: PrimitiveAtom<Value>,
  reducer: (v: Value, a: Action) => Value,
): [Value, (action: Action) => void]

尚、useReducerAtom の型定義に出てくる PrimitiveAtom というのはプリミティブ値の Atom ではないのでオブジェクト型やユニオン型の Atom も対象にできます。具体的には atom(0) で作成した Atom のような「値の取得」と「 Value | (prev: Value) => Value の形式で値の更新」ができる Atom の事で、Read-only Atom や Write-only Atom のような 派生 Atom が対象外となります。

指定した Atom のみ Store を分離する

Scope — Jotai, primitive and flexible state management for React
Provider と Store の説明でも触れた通り通常は Provider のツリー毎に Store が分離されるのですが、アプリケーションによってはグローバルに扱いたいものの方が多くて Store を分離したい Atom は一部だけの場合もあります。このケースの場合だと useAtom などで都度 { store } のオプションを指定するのは冗長になってしまいます。これを解決してくれるのが拡張機能の jotai-scope から提供されている ScopeProvider です。

利用するにはまずパッケージにインストールが必要です。

pnpm add jotai-scope

ScopeProvider の利用は簡単で props の atoms に Store を分離したい atom config を配列で渡すだけです。atoms に指定された Atom だけが ScopeProvider の子コンポーネントに対して分離された値を持つようになります。

import { ScopeProvider } from 'jotai-scope'

const ScopeCounter = () => (
  <ScopeProvider atoms={[countAtom]}>
    <Counter />
  </ScopeProvider>
)

Atom の一部のみを参照する

Optics — Jotai, primitive and flexible state management for React
Record を含むプロパティの多いオブジェクト型の Atom を参照する時にコンポーネントによってはオブジェクトの一部だけが必要で、オブジェクトの他の部分が更新された時に再描画されたくないケースがあります。これを解決してくれるのが拡張機能の jotai-scope から提供されている focusAtom です。

pnpm add optics-ts jotai-scope

import { atom, useAtom } from 'jotai'
import { focusAtom } from 'jotai-optics'

const baseAtom = atom({ a: 1, b: 2 })
const subAtom = focusAtom(baseAtom, (optic) => optic.prop('a'))

const Example = () => {
  const [a, setA] = useAtom(subAtom)
}

focusAtom は第一引数に atom config、第二引数に optic を受け取って optic.prop('key') の形式で参照したいプロパティ名を指定して返す関数を渡します。focusAtom はこの第二引数で指定したプロパティのみを参照する派生 Atom を作成してくれます。

focusAtom で作成した atom config は他と同様に useAtom などで利用できます。値を参照時は派生元のオブジェクトの他のプロパティが更新されても再描画が発生しませんし、値を更新時は派生元の対象プロパティの値が更新されます。

下記のコードはカスタムフック内で focusAtom を利用した応用的な例です。

import { atom } from 'jotai'
import { focusAtom } from 'jotai-optics'

type PriceRecord = Record<string, number>
const priceRecordAtom = atom<PriceRecord>({})

const useItemPrice = (name: string) => {
  // 通常の Atom と同奥にメモ化は必須
  const priceAtom = useMemo(() => focusAtom(
    priceRecordAtom,
    (optic) => optic.prop(name)
  ), [name])
  
  return useAtom(priceAtom)
}

optics-ts

Method chaining API - optics-ts
具体的にどういった事ができるのかまでは把握していないのですが、focusAtom は上記の optics-ts リファレンスにある Isomorphisms, Lenses, Prisms の 3 つグループにあるメソッドの返す値をサポートしているようです。

atomFamily

Family — Jotai, primitive and flexible state management for React
Recoil から Jotai に移行する際に Recoil の atomFamily の代替として振る舞いなどを調べたのですが、個人的な結論はなるべく利用を避けて atom 関数を利用した実装の方がいいと感じました。せっかく調べたというのもありますので atomFamily の振る舞いと避けた方がいい理由について説明します。

atomFamily は辞書形式の Atom で与えられたパラメータに一致する Atom を返します。jotai/utils から提供されている atomFamily を使って atom config を作成するのですが、この型定義は以下のようになっています。

export function atomFamily<Param, AtomType extends Atom<unknown>>(
  initializeAtom: (param: Param) => AtomType,
  areEqual?: (a: Param, b: Param) => boolean,
): AtomFamily<Param, AtomType>

initializeAtom は atomFamily 内の Atom を参照する時に渡すパラメータ param を受け取り、初期値となる Atom を返す関数を指定します。Atom の初期値にパラメータを利用したいケースを想定して受け取れるようになっていますが、必要なければ省略可能です。

areEqual は atomFamily を利用する際に受け取ったパラメータと内部のパラメータが一致するかに利用する為の関数を指定します。例えばオブジェクトの id プロパティだけで一致するとみなしたい時は (a, b) => a.id === b.id のような関数を渡します。省略可能でデフォルトは Object.is が使用されます。

下記は atomFamily の型パラメータも指定したサンプルコードです。

import { atom, type Atom, useAtom } from 'jotai'
import { atomFamily } from 'jotai/utils'

type UserId = number
type User = { id: UserId; name: string}

const userFamily = atomFamily<UserId, Atom<User | undefined>(
  () => atom(undefined)
)

const Example = () => {
  // UserId: 1 の User データを参照、初期値は上で指定した undefined
  const [user, setUser] = useAtom(userFamily(1))
}

次に内部実装に関連した扱う際の注意点について説明します。

通常の Atom は内部で WeakMap を利用した実装になっているので参照されなくなったら GC されるのですが、atomFamily は内部で Map を利用しており参照されなくなっても GC されません。つまりグローバル変数に Map を持つことと同意で、提供されている remove や setShouldRemove のメソッドを使って管理する必要があります。

このような実装になっているので、アプリケーションと生存期間が同じで全ての値を常に保持する必要があるケースや、保持する値が限られているケースでは問題ありませんが、そうでない場合は注意が必要です。

add builtin weakmap version of atomFamily? · pmndrs/jotai · Discussion #2239
このあたりは作者も Discussions で触れられており、Jotai のコンセプトとは異なる存在なので、将来的には jotai/utils から拡張機能などに切り出す事を考えてるようです。