TypeScript for atom with setter whose update type is different than the value type

[elliotwaite]

I have two Theme objects, lightTheme and darkTheme, and I want an atom whose value will be one of them, but I want to be able to set the atom by passing in a string of either "light" or "dark". Getting the types to work was tricky, but I came up with a solution, I was just wondering if there was a simpler way of doing this. Here's my current solution (demo):

const lightTheme = { backgroundColor: "#ffffff" };
const darkTheme = { backgroundColor: "#303030" };

type Theme = typeof lightTheme;
type LightOrDark = "light" | "dark";

const themeAtom: WritableAtom<Theme, LightOrDark> = atom(
  lightTheme,
  (_get, set, arg) =>
    set(
      (themeAtom as unknown) as WritableAtom<Theme, Theme>,
      arg === "light" ? lightTheme : darkTheme
    )
);

It seems like I had to combine both WritableAtom<Theme, LightOrDark> and WritableAtom<Theme, Theme> to get it to work.

Here are some of the other solutions I tried that didn't work. When I tried it without types:

const themeAtom = atom(
  lightTheme,
  (get, set, arg) => set(themeAtom, arg === "light" ? lightTheme : darkTheme)
);

... I got this error: 'themeAtom' implicitly has type 'any' because it does not have a type annotation and is referenced directly or indirectly in its own initializer.

When I tried it with just WritableAtom<Theme, LightOrDark>:

const themeAtom: WritableAtom<Theme, LightOrDark> = atom(
  lightTheme,
  (get, set, arg) => set(themeAtom, arg === "light" ? lightTheme : darkTheme)
);

... I got an error on the arg === "light" ? lightTheme : darkTheme part that said: Argument of type '{ backgroundColor: string; }' is not assignable to parameter of type 'LightOrDark'. Type '{ backgroundColor: string; }' is not assignable to type '"dark"'.

When I tried it with just WritableAtom<Theme, Theme>:

const themeAtom: WritableAtom<Theme, Theme> = atom(
  lightTheme,
  (get, set, arg) => set(themeAtom, arg === "light" ? lightTheme : darkTheme)
);

... I got an error on the arg === "light" part that said: This condition will always return 'false' since the types '{ backgroundColor: string; }' and 'string' have no overlap.

I know that I could also use a separate base atom and a derived atom (which works without having to specify types):

const baseThemeAtom = atom(lightTheme);
const themeAtom = atom(
  (get) => get(baseThemeAtom),
  (get, set, arg) =>
    set(baseThemeAtom, arg === "light" ? lightTheme : darkTheme)
);

... but I was hoping to find a simple way to get it to work using just a single atom.

Note: All of the versions listed above are also in the demo.

[dai-shi]

This is sort of typescript limitation, and that's why we don't provide such examples in our docs.
In JS, it's valid. And, in TS, your solution is fine too.
Our current recommendation is to use two atoms, or use atomWithReducer (implemented with any).

[elliotwaite]

Okay, sounds good. Just wanted to make sure I wasn't missing a simpler solution. Thanks for the answer.

[elliotwaite]

I guess one of the limiting factors is that TypeScript currently lacks the ability to check for reference equality, since we would need something like Object.is() in TypeScript to be able to check if the Setter was being called from within the Write function of an atom, and if that atom was the same atom being passed into the Setter, in which case the Setter's value argument should be of type Value, otherwise it should be of type Update. Hopefully, TypeScript will add something like this in the future.

[quangbuule]

This is mindblowing, I didn't know that you can set the atom in write fn without causing an infinite looping. For this, I would suggest:

const lightTheme = { backgroundColor: "#ffffff" };
const darkTheme = { backgroundColor: "#303030" };

type Theme = typeof lightTheme;
type LightOrDark = "light" | "dark";

const themeAtom: WritableAtom<Theme, LightOrDark | Theme> = atom(
  lightTheme,
  (_get, set, arg) => {
    if (arg === 'light') {
      set( themeAtom, lightTheme);
    }

    if (arg === 'dark') {
      set(themeAtom, darkTheme);
    }
  }
);

[elliotwaite]

This is another option, however, this reduces type safety since we only want to be able to call setThemeAtom(_) with a value of type LightOrDark, and then from within themeAtom's write functions, we only want to be able to call set(themeAtom, _) with a value of type Theme. If we set the type to WritableAtom<Theme, LightOrDark | Theme> then both of those calls will accept a value of type LightOrDark | Theme.

This is mindblowing, I didn't know that you can set the atom in write fn without causing an infinite looping.

Yeah, it's a nice option to know about. I didn't realize this was possible until I read the source code. The setter that gets passed into the atom's write function will check if the atom being set with it is the same atom that the write function is for, and if so, it will set the value directly instead of recursively calling its write function again. Here is the code (I replaced the none relevant parts with ..., see here for the full code):

    const setter: Setter = <V, U, R extends void | Promise<void>>(
      a: WritableAtom<V, U, R>,
      v?: V
    ) => {
      ...
      if ((a as AnyWritableAtom) === atom) {
        ...
        // This sets the value directly.
        setAtomPromiseOrValue(version, a, v)
        ...
      } else {
        // This calls the `a` atom's update function with the value.
        promiseOrVoid = writeAtomState(version, a as AnyWritableAtom, v)
      }
      ...
    }