Reading Code - Chakra UI

I spent over a year as a principal engineer in one of News UK’s platform teams where we were building libraries and tooling for the other engineers in the company. One of those tools was a fantastic design system.

And even though I had nothing to do with it, it got me to think more deeply about reusable components.

The most popular libraries in the React community have feature parity and differ primarily in their components’ APIs. So deciding on one is more a matter of taste and visual preference than a matter of technical advantage.

I decided to stroll through Chakra UI’s codebase because I’ve heard overwhelmingly positive things about it, yet I’ve never got the chance to use it. This is my chance to get familiar with it and, hopefully, learn something interesting while I’m doing so.

The Approach

I’ve written other articles where I read open-source projects’ code, but before we dive into Chakra, let’s make a quick refresher on how we will do that. The approach I usually take when exploring an unfamiliar codebase is to pick one of its public APIs and follow the execution path all the way through.

In Chakra’s context, this will be one of its components. I settled on a button because it’s simple enough to understand how it works but still has some small amount of functionality attached to it.

The Repository

We see immediately that Chakra is using a monorepo, and every component is implemented as its own package. This gives users the nice ability to only install the separate components we need instead of pulling the entire library.

In some cases, this may be a worthy trade-off.

But besides the components, Chakra’s utility functions and hooks are also implemented as separate packages. Something that caught my eye is the presence of a core folder which tells me that the components probably rely on a common set of functionality a la react-query.

├── packages
|   ├── components
|   ├── core
|   ├── hooks
|   ├── utilities
|   ├── ...

Separating the core functionality from the bindings is a practice I keep seeing in the community. It opens the door for easier portability between frameworks in the future.

The Button

The button is a self-contained package that holds all related functionality, together with tests and stories. The src folder holds all the actual logic (following the common convention), well-split into multiple files.

├── stories
├── tests
├── src
|   ├── button.tsx
|   ├── button-context.tsx
|   ├── button-group.tsx
|   ├── button-icon.tsx
|   ├── button-spinner.tsx
|   ├── button-types.tsx
|   ├── icon-button.tsx
|   ├── index.tsx
|   ├── use-button-type.tsx
├── package.json
├── tsconfig.json
├── docs.json
├── README.md

Having an index.ts file for modules and folders is a good practice, so you can export everything you want to make public in a single place. Of course, nothing is stopping people from importing them directly from the file, but this defines what is made to be public, and provides shorter import paths as a side-effect.

Forwarding Refs

When we open button.tsx, we will immediately notice that the component is not a simple function. It’s wrapped in another one called forwardRef.

export const Button = forwardRef<ButtonProps, 'button'>(
  (props, ref) => {
    // ...
  }
)

After following the forwardRef function we see that it’s a wrapper around React’s built-in forwardRef but adds Chakra specific types. Notice how the function is named forwardReactRef to make sure there is no obscurity.

export function forwardRef<
  Props extends object,
  Component extends As
>(
  component: React.ForwardRefRenderFunction<
    any,
    RightJoinProps<PropsOf<Component>, Props> & {
      as?: As
    }
  >
) {
  return forwardReactRef(component) as unknown as ComponentWithAs<
    Component,
    Props
  >
}

Chakra’s forwardRef accepts generic types to avoid type casting outside of it, and it also sets the returned type from React’s built-in function, named forwardReactRef here. For some extra context, forwarding a ref is a technique to pass down a ref passed to the main component.

So a ref passed to Chakra’s Button will be passed down in case the user wants to have a reference to the underlying DOM node. This is a practice often used in component libraries.

An interesting detail is the type casting as unknown as ComponentWithAs.

This is an interesting detail that I’ve seen in a couple of codebases and thought to be an error before digging into it here.

The technique is known as double assertion and is used to make sure that the type casting succeeds. Casting A to B will only work if A can actually be assigned to B or B to A.

By casting it to unknown first, we ensure that there won’t be problems with the types.

Contexts

The first line of code in the component’s definition is a call to a custom hook.

export const Button = forwardRef<ButtonProps, 'button'>(
  (props, ref) => {
    const group = useButtonGroup()
    // ...
  }
)

I’m a big advocate of using custom hooks as an idiomatic form of abstraction in React and opened the useButtonGroup implementation with great interest.

export const [ButtonGroupProvider, useButtonGroup] =
  createContext<ButtonGroupContext>({
    strict: false,
    name: 'ButtonGroupContext',
  })

You could argue that the custom hook is unnecessary because it’s hiding a single call to createContext underneath. But it’s not about the lines of code that it’s saving, but the abstraction it is providing. By using it, the Button component doesn’t have to be aware of the context underneath.

But you will notice that the returned value from createContext is not what you would expect when creating a new context in React. This is because it is once again a wrapped function. In fact, this is a pattern that we will continue seeing.

It’s something I’ve written about in my book Tao of React - it’s good to abstract functions and components, especially 3rd party ones, so you can have better control over their API.

export function createContext<T>(options: CreateContextOptions<T> = {}) {
  // ...
  const Context = createReactContext<T | undefined>(undefined)
  // ...

  function useContext() {
    const context = useReactContext(Context)

    if (!context && strict) {
      const error = new Error(
        errorMessage ?? getErrorMessage(hookName, providerName),
      )
      // ...
      throw error
    }

    return context
  }

  return [Context.Provider, useContext, Context] as CreateContextReturn<T>
}

I’ve omitted some details for brevity, but the main idea is well illustrated here. By wrapping the native context creation function, the authors can standardize the context structure by specifying the type of the options argument.

They still accept a generic CreateContextOptions<T> to define specific fields.

They also create their own wrapper around the native useContext and add an additional check if the context was marked as strict. Trying to use one such context outside of its boundaries will result in an error to avoid potentially bigger problems.

Styles and Abstractions

Going back to the button component, on the next line we see a call to another custom hook.

export const Button = forwardRef<ButtonProps, "button">((props, ref) => {
  const group = useButtonGroup()
  const styles = useStyleConfig("Button", { ...group, ...props })
  // ...
})

I was curious how styles are set up here so I followed the custom hook to an implementation that seemed familiar.

export function useStyleConfig(
  themeKey: string,
  props: ThemingProps & Dict = {},
) {
  return useStyleConfigImpl(themeKey, props) as SystemStyleObject
}

The custom hook acts only as a public wrapper to an internal hook that is not exported. I first saw this convention in react-query’s codebase, but I thought it was something that its authors came up with.

After some digging and asking around on social media, it turned out that this is a common pattern in Java. Since the JS ecosystem is not great at defining what idiomatic code looks like, many of the people coming from other languages are replicating techniques they know.

This is a way to decouple an interface from an implementation.

But since the function is not doing any additional checks or modifying the internal function’s API, I don’t see much benefit of having the extra wrapper. The useStyleConfigImpl itself is not that interesting, but there’s one bit that I want to highlight.

function useStyleConfigImpl(
  themeKey: string | null,
  props: ThemingProps & Dict = {},
) {
  // ...

  /**
   * Store the computed styles in a `ref` to avoid unneeded re-computation
   */
  const stylesRef = useRef<StylesRef>({})

  // ...

  return stylesRef.current
}

The comment above the stylesRef variable is very insightful and a good example of context that can’t be conveyed only through code. It doesn’t describe how something is done, but why.

Memoization

The next thing that I want to highlight in the component is that they’ve used useMemo to memoize the computed styles object for the component.

export const Button = forwardRef<ButtonProps, "button">((props, ref) => {
  // ...

  const buttonStyles: SystemStyleObject = useMemo(() => {
    const _focus = { ...styles?.["_focus"], zIndex: 1 }
    return {
      display: "inline-flex",
      appearance: "none",
      alignItems: "center",
      justifyContent: "center",
      userSelect: "none",
      position: "relative",
      whiteSpace: "nowrap",
      verticalAlign: "middle",
      outline: "none",
      ...styles,
      ...(!!group && { _focus }),
    }
  }, [styles, group])

  // ...
})

There are a lot of conflicting opinions in the React community about when and how useMemo should be used. It’s often seen as a premature optimization, and I’m leaning toward this opinion as well.

According to benchmarks, useMemo’s cache lookup is slower than just doing the data modification yourself for looping through a hundred-item list, for example. But using useMemo is not only about speed but preventing re-renders.

In fact, there are companies that put useMemo and useCallback around every computed value and function they use to prevent unnecessary rendering. This approach follows the philosophy that a guaranteed but small performance hit is preferable to a potentially bigger problem.

Rendering

Not we’ve reached the return statement of the component, and things get even more interesting here. You would expect a simple HTML <button> tag underneath, but it turns out that the component we’ve been reviewing so far is only a layer.

export const Button = forwardRef<ButtonProps, "button">((props, ref) => {
  // ...
  return (
    <chakra.button
      ref={useMergeRefs(ref, _ref)}
      as={as}
      type={type ?? defaultType}
      data-active={dataAttr(isActive)}
      data-loading={dataAttr(isLoading)}
      __css={buttonStyles}
      className={cx("chakra-button", className)}
      {...rest}
      disabled={isDisabled || isLoading}
    >
      ...
    </chakra.button>
  )
})

The returned JSX is not HTML but refers to <chakra.button> and other elements from an underlying styled system.

In fact, this styled system is available directly as a separate package, so people who want to take advantage of the Chakra API without using the actual components can do so.

This could be a good base to develop a design system of your own.

Additionally, you can call the chakra helper with your own component to wrap it and make its props consistent with the rest of your components.

One of the rules in my book about React focuses on the importance of consistency. To reduce the complexity of your application, you should have all your components follow the same design principles, and this factory function helps you achieve it.

The Chakra Factory

The chakra factory is a very interesting design decision. It gives the foundation on top of which the Chakra components are built. So in effect, it’s a low-level component library that extends the native HTML elements.

function factory() {
  const cache = new Map<DOMElements, ChakraComponent<DOMElements>>()

  return new Proxy(styled, {
    /**
     * @example
     * const Div = chakra("div")
     * const WithChakra = chakra(AnotherComponent)
     */
    apply(target, thisArg, argArray: [DOMElements, ChakraStyledOptions]) {
      return styled(...argArray)
    },
    /**
     * @example
     * <chakra.div />
     */
    get(_, element: DOMElements) {
      if (!cache.has(element)) {
        cache.set(element, styled(element))
      }
      return cache.get(element)
    },
  }) as ChakraFactory & HTMLChakraComponents
}

export const chakra = factory()

The implementation is also fascinating. The chakra object is a Proxy to another object called styled. Upon finding this I had to do a quick search and remember what a Proxy even was, because I’ve never had to use one so far.

It’s an object that can be used in place of the original object, but it gives you the ability to redefine some of its methods or behavior.

In the context of Chakra UI, this proxy helps it provide a utility that can be used both as an object and as a factory function.

If chakra is called as a function, the apply method will be called and it will return a wrapped styled component. But if it’s called as an object, it will create a new element and cache it to improve performance.

See the example comments in the code above for more details.

By using a Proxy, the authors provide a more flexible API to the users, and hide the details of how the elements are created underneath. It’s definitely an interesting decision, aimed at pushing more complexity behind an abstraction.

The Styled Function

To get to the end of the execution flow, we have to follow the call to styled to see what’s hiding in there. Knowing how many layers we went through so far, it’s no surprise that styled is yet another wrapper on top of a different function.

export function styled<T extends As, P = {}>(
  component: T,
  options?: ChakraStyledOptions,
) {
  // ...
  const Component = emotion_styled(
    component as React.ComponentType<any>,
    styledOptions,
  )(styleObject)

  const chakraComponent = React.forwardRef(function ChakraComponent(
    props,
    ref,
  ) {
    // ...
    return React.createElement(Component, {
      ref,
      "data-theme": forced ? colorMode : undefined,
      ...props,
    })
  })

  return chakraComponent as ChakraComponent<T, P>
}

Underneath, it makes a call to Emotion, a popular CSS-in-JS library, and then manually creates the component we need by calling React.createElement. And with that, we’ve reached the end, and have a pretty good idea about how a component is created in Chakra UI.

The Main Takeaway

The most important thing to carry home is the importance of layers. Too often, we overlook this and bundle too much logic into our React components.

This button could’ve been developed by making a call to Emotion directly. But the added layers give flexibility and decouple the codebase from the libraries they depend on. Wrapping the functions could’ve been avoided. But doing so enforces standards and consistency across all components in the library.

You needn’t apply all these principles in your React codebase starting tomorrow. However, they are a good inspiration of how you can make your codebase more decoupled and well-structured. The entire Chakra library depends on Emotion, but it’s so well abstracted that you can work without noticing it.