Why Jotai Is the Key to Better State Management in React

Why Jotai Is the Key to Better State Management in React

If I started talking to you about state management with React, the first thing that may come to mind is Context. Or maybe react-redux . For what seemed like an eternity, React's Context and Redux were used as the main solutions for providing data used by many components within applications.

React's built-in state solution may seem to be the most viable solution in the short term, but can cause maintenance headaches and performance issues down the line. Let me explain.

React Context API

First, we'll look at how to use Context with this simple dark mode toggle.

So, let's begin by creating a context.

import React from "react";

export type Theme = "light" | "dark";

export const ThemeContext = React.createContext<{
    theme: Theme;
    setTheme: React.Dispatch<React.SetStateAction<Theme>>;
}>({ theme: "light", setTheme: () => {} });

This exported ThemeContext is later used to initialise the context value, as well as retrieve its value within other components.

Note that React doesn't provide us with a direct method of updating the context data. That is why we must also have setTheme to update the value within the context.

For example, let's create a button to toggle the theme.

Consuming context

In React function components, context is consumed using the useContext hook. Class components achieve the same effect by setting MyClass.contextType = ThemeContext (for example).

import React, { useContext } from "react";
import { ThemeContext } from "./store";
import "./ThemeToggle.css";
import { Sun, Moon } from "react-feather";

export default function ThemeToggle() {
    const context = useContext(ThemeContext);

    const onToggle = () => {
        const newTheme = context.theme === "dark" ? "light" : "dark";
        context.setTheme(newTheme);
    };

    return (
        <button type="button" className={context.theme} onClick={onToggle}>
            {context.theme === "dark" ? (
                <Sun className="toggle sun" />
            ) : (
                <Moon className="toggle moon" />
            )}
        </button>
    );
}

When it's dark, we want to show the sun to know that clicking the button makes the screen light. And when it's light, we display the moon to shift to the dark side when clicked.

Providing context data

Finally, when passing down data from a context, we wrap components with a Provider . Typically, the value used in a context provider comes from state variables.

import { useState } from "react";
import "./App.css";
import ThemeToggle from "./ThemeToggle";
import { Theme, ThemeContext } from "./store";
import ThemeStatus from "./ThemeStatus";

function App() {
    const [theme, setTheme] = useState<Theme>("light");

    return (
        <ThemeContext.Provider value={{ theme, setTheme }}>
            <div className={`App ${theme}`}>
                <h1>Theme Mode Demo</h1>
                <ThemeToggle />
                <ThemeStatus />
            </div>
        </ThemeContext.Provider>
    );
}

export default App;

And that would be that.

Overview

So, now we've seen how easy it is to implement Context in your React project. Looks simple, right? Well, for small applications where not much state is needed, this solution is fine.

The problem arises when you're application starts to become much more complex. Imagine you're building a multi-user chat application. Now instead of just a dark mode toggle, you also have to manage the currently authenticated user, the locale to translate the site content to, the user's site appearance preferences, and so on.

One way of going around this would be to use just one global context to store everything. This way, only one context provider is created within the application. However, if lots of components depend on the same context, any change made to the value will trigger every component to re-render. As mentioned before, this can cause major performance issues as the app keeps growing.

The alternative would be to create multiple contexts for different purposes. For example, a context to store user data, a context to store the language code, a context to store the user's appearance settings, and so on. Not even to mention if we wanted to persist any of our variables in localStorage. The issue now is that we create all these different contexts, all of which need to be initialised, all of which need slightly different code to consume. And all we wanted initially was a dark mode.

This also makes isolated unit tests even harder. We now have to figure out the context consumed by the components we test and make sure they are mocked correctly.

As you can see, React Context API can be used minimally, but there are still many unsustainable ways to use it. Ideally, we should try to not rely on the global state and instead try to avoid it if data doesn't need to be passed so deeply.

Now, what if there was a way we could reap the benefits of a simple-to-use API whilst avoiding unnecessary re-renders? What if there was a library designed to make state management flexible yet simple, whilst avoiding the issues described above and integrating extra features such as storage persistence and state derivation?

Meet Jotai.

Jotai

Jotai solves all of our problems whilst still being scalable to both large and small projects. Install it with npm i jotai and let's get into another example.

This time, we'll build a simple word counter like below.

What is an atom?

Each piece of state in Jotai is known as an "atom". Just like React Context , we export configurations for atoms to use elsewhere.

import { atom } from "jotai";

// Initialise the value to an empty string
export const textAtom = atom("");

Note that all atoms are made available globally by default, so we won't even need to include a single Provider in our example.

Jotai also makes deriving state easy. We can create a new atom which derives the textAtom to calculate the number of characters of text, for example. The get function passed in the callback allows us to read atom values whilst tracking the new atom as a dependency.

export const charactersAtom = atom((get) => {
    return get(textAtom).length;
});

Consuming atoms with Jotai is designed to be as simple as calling React.useState. We can create a component to read input from the user with the following code.

import React from "react";
import { useAtom } from "jotai";
import { textAtom } from "./store";
import "./TextInput.css";

export default function TextInput() {
    const [text, setText] = useAtom(textAtom);

    const onChange = (event: React.ChangeEvent<HTMLTextAreaElement>) => {
        setText(event.target.value);
    }

    return <div>
        <textarea title="Enter text" value={text} onChange={onChange} />
    </div>
}

As well as useAtom , Jotai also provides us with useAtomValue and useSetAtom if we only want to read or update an atom value.

So under the <textarea />, we can show a display of information about the text entered. Here are just a few more derived atoms.

import { atom } from "jotai";
import { atomWithStorage, atomWithReset } from "jotai/utils";

// Matches any sentence terminating character
const STOP_WORDS = /[\.?!]/g;

const splitText = (text: string, splitter: string | RegExp) => {
    return text.split(splitter).filter((part) => part !== "");
}

export const textAtom = atomWithReset("");
export const sentencesAtom = atom((get) => {
    return splitText(get(textAtom), STOP_WORDS).length;
});

export const wordsAtom = atom((get) => {
    return splitText(get(textAtom), " ").length;
});

export const charactersAtom = atom((get) => {
    return get(textAtom).length;
});

export const paragraphsAtom = atom((get) => {
    return splitText(get(textAtom), "\n").length;
});

// Calculated using the Coleman-Liau index
// See https://en.wikipedia.org/wiki/Coleman-Liau_index
export const readabilityAtom = atom((get) => {
    const words = get(wordsAtom);

    if (words < 150) {
        return NaN;
    }

    const characters = get(charactersAtom);
    const sentences = get(sentencesAtom);

    const L = (characters / words) * 100;
    const S = (sentences / words) * 100;

    const readability = 0.0588 * L - 0.296 * S - 15.8;
    return readability.toFixed(1);
});

And now we can display all of the values.

import React from "react";
import { useAtomValue } from "jotai";
import { wordsAtom, charactersAtom, sentencesAtom, paragraphsAtom, readabilityAtom } from "./store";
import "./TextStats.css";

export default function TextStats() {
    const words = useAtomValue(wordsAtom);
    const characters = useAtomValue(charactersAtom);
    const sentences = useAtomValue(sentencesAtom);
    const paragraphs = useAtomValue(paragraphsAtom);
    const readability = useAtomValue(readabilityAtom);

    return <div className="stats">
        <div>
            <h4>Words</h4>
            <span>{words}</span>
        </div>
        <div>
            <h4>Characters</h4>
            <span>{characters}</span>
        </div>
        <div>
            <h4>Sentences</h4>
            <span>{sentences}</span>
        </div>
        <div>
            <h4>Paragraphs</h4>
            <span>{paragraphs}</span>
        </div>
        <div>
            <h4>Readability</h4>
            <h6>(&gt; 150 words)</h6>
            <span>{readability}</span>
        </div>
    </div>
}

Actions

Action atoms

Now we've learnt how to create readable derived atoms, let's talk about action atoms.

Action atoms are created with a similar pattern, but only exist to write a value to an atom. In the context of our example, we can create an action atom to convert the entered text into a readable case.

export const sentenceCaseAtom = atom(null, (_get, set) => {
    set(textAtom, (prev) => {
        // Converts the previous value of the atom into sentence case.
        return prev.split(" ").map(word => {
            return word[0].toUpperCase() + word.slice(1);
        }).join(" ")
    })
});

Resettable atoms

Another cool feature of Jotai (specifically Jotai's utility bundle) is atoms that can have their values reset. So instead of having to use useSetAtom along with a decided default value, you can use useResetAtom to provide a function that only does a reset.

import { atomWithReset } from "jotai/utils";
export const textAtom = atomWithReset("");

Usage

import React from "react";
import { useResetAtom } from "jotai/utils";
import { useSetAtom } from "jotai";
import { sentenceCaseAtom, textAtom } from "./store";
import "./TextActions.css";

export default function TextReset() {
    const resetText = useResetAtom(textAtom);
    const toSentenceCase = useSetAtom(sentenceCaseAtom);

    return <div className="actions">
        <button type="button" onClick={resetText}>Reset</button>
        <button type="button" onClick={toSentenceCase}>Format case</button>
    </div>
}

Summary

As you've seen, Jotai is a great, minimalistic state management solution for React. Currently, I use it in almost all my new React projects instead of useContext or other tools like Redux. And I think you should consider doing the same in your future projects.

Despite the simplistic core API, Jotai's utilities provide a lot more features than React does out-of-the-box. Furthermore, the primitive nature of their atom-driven state is much more maintainable than Context as it avoids having many Provider with many state variables, which ends up becoming hard to read and therefore hard to test.

Just one thing to note: make sure to separate your atom configurations with Jotai by keeping related atoms in the same file.

If you enjoyed reading this article, consider following me for more.

Did you find this article valuable?

Support Wool Doughnut by becoming a sponsor. Any amount is appreciated!