Implementing Dark Mode with React, Jotai, and Custom Hooks

Note: I've been programming as a hobby for a while and as the twilight of my Army career finally approaches, I've been working to learn modern web development tools for a career when I get out of the military. My current project is a task manager and I've pulled the code from there for this article.

Dark mode has become a staple feature in modern web applications, offering users a more comfortable viewing experience in low-light environments. In this article, we'll explore how to implement a robust dark mode toggle using React, Jotai for state management, and a custom hook for better organization and reusability.

The Building Blocks

Let's break down the key components of our implementation:

1. Jotai for state management
2. A custom hook for UI-related operations
3. A DarkModeToggle component

1. Jotai for State Management

I'm using Jotai, an atomic state management library for React. It came as a recommendation from a friend in web development when I was complaining about how boilerplatey the flux pattern was. After telling me people weren't really hot on using flux (Redux) anymore, he pointed me towards Jotai or MobX. Jotai made the most sense to me for this project so I settled on it and the atomic approach to state. Here's how I set up our UI state:

// src/store/uiAtoms.ts
import { atom } from "jotai"
import { atomWithStorage } from "jotai/utils"
import type { ThemeMode, UIState } from "@/types"

export const uiStateAtom = atomWithStorage<UIState>("uiState", {
  darkMode: "system",
  showCompleted: false, //this line and down are for other parts of my UI state
  showOverdue: false,
  // ... truncated to keep the shared code down
})

export const setDarkModeAtom = atom(null, (get, set, mode: ThemeMode) => {
  const currentState = get(uiStateAtom)
  set(uiStateAtom, { ...currentState, darkMode: mode })
})

I'm using atomWithStorage to persist the UI state in local storage, ensuring that the user's preference is remembered across sessions. I had written some code of my own to do this before stumbling on the fact that Jotai already did it for me - the story of my life.

2. Custom Hook for UI Operations

To encapsulate UI-related logic, we've created a custom hook:

// src/hooks/useUI.ts
import { useAtom } from "jotai"
import { uiStateAtom, setDarkModeAtom } from "@/store/uiAtoms"

export const useUI = () => {
  const [uiState] = useAtom(uiStateAtom)
  const [, setDarkMode] = useAtom(setDarkModeAtom)

  return {
    uiState,
    setDarkMode,
  }
}

The hook provides a clean interface for components to interact with the UI state.

3. DarkModeToggle Component

Now, let's look at the DarkModeToggle component. I'm utilizing shadcn's wonderful DropdownMenu component after realizing it's built on radix (the library I was looking at using because it takes care of all the accessibility stuff for me).

// src/components/common/DarkModeToggle.tsx
import React, { useEffect } from "react";
import { useUI } from "@/hooks/useUI";
import { Moon, Sun } from "lucide-react";
import { Button } from "@/components/ui/button";
import {
  DropdownMenu,
  DropdownMenuContent,
  DropdownMenuItem,
  DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";

const DarkModeToggle: React.FC<{ className?: string }> = ({ className }) => {
  const { setDarkMode, uiState } = useUI();

  useEffect(() => {
    const root = window.document.documentElement;
    const prefersDark = window.matchMedia(
      "(prefers-color-scheme: dark)"
    ).matches;

    if (
      uiState.darkMode === "dark" ||
      (uiState.darkMode === "system" && prefersDark)
    ) {
      root.classList.add("dark");
    } else {
      root.classList.remove("dark");
    }
  }, [uiState.darkMode]);

  return (
    <DropdownMenu>
      <DropdownMenuTrigger asChild>
        <Button
          variant="outline"
          size="icon"
          className={className}>
          <Sun className="h-4 w-4 rotate-0 scale-100 transition-all dark:-rotate-90 dark:scale-0" />
          <Moon className="absolute h-4 w-4 rotate-90 scale-0 transition-all dark:rotate-0 dark:scale-100" />
          <span className="sr-only">Toggle theme</span>
        </Button>
      </DropdownMenuTrigger>
      <DropdownMenuContent align="end">
        <DropdownMenuItem onClick={() => setDarkMode("light")}>
          Light
        </DropdownMenuItem>
        <DropdownMenuItem onClick={() => setDarkMode("dark")}>
          Dark
        </DropdownMenuItem>
        <DropdownMenuItem onClick={() => setDarkMode("system")}>
          System
        </DropdownMenuItem>
      </DropdownMenuContent>
    </DropdownMenu>
  );
};

export default DarkModeToggle;

This component uses the useUI hook to access and modify the dark mode state. It also includes a useEffect hook to apply the appropriate class to the document root based on the current mode.

Putting It All Together

With these pieces in place, implementing dark mode becomes straightforward:

1. The uiStateAtom stores the current dark mode preference.
2. The useUI hook provides an easy way to access and modify this state.
3. The DarkModeToggle component allows users to change their preference.
4. The useEffect hook in DarkModeToggle applies the appropriate class to enable/disable dark mode.

This approach offers several benefits:

• Persistence: User preferences are stored in local storage.
• System Preference Support: The implementation respects the user's system-wide dark mode setting.
• Separation of Concerns: State management is handled by Jotai, while component logic is encapsulated in a custom hook.
• Reusability: The useUI hook can be easily extended to handle other UI-related state.

Conclusion

The code above is a pretty contrived example of using custom hooks, but dark mode really is something that everyone should know how to implement, newbies like myself included. It's one of the first things I look at how to implement when messing around with different libraries and frameworks.

I've really fallen in love with using custom hooks for separating state/stores from the logic needed for components. It's probably one of my favorite patterns I've discovered thanks to (finally) getting around to learning React. My skills have a long way to go (and constructive feedback - both positive and negative - are very welcome)