[Proposal] atomWithWrapper

[TwistedMinda]

Hello,

I'm currently trying to setup a clean opinionated approach for using jotai, cause I'm finding approximately none in the wild. Planning to make the big step & start using it in production ASAP.
My constraints are the following:

• separate groups of logic in "modules"
• keep component/modules imports untied to underlying library
• typescript must be respected all the way up

I've come up with 2 main elements:

• A "State" wrapper for my exposed atoms :

type Writable<Value> = WritableAtom<Value, any, any>
type Store<Value, FamilyParam = any> = Atom<Value> | Writable<Value> | AtomFamily<FamilyParam, Atom<Value>>

export class State<Value, FamilyParam> {
  reference: Store<Value>

  constructor(reference: Store<Value, FamilyParam>) {
    this.reference = reference
  }

  ref(param?: any) {
    if (typeof this.reference === 'function')
      return (this.reference as AtomFamily<FamilyParam, Atom<Value>>)(param) // It was an atom family
    return this.reference
  }
  use(param?: any) { return useAtom(this.ref(param) as Atom<Value>) }
  get(param?: any): Value { return useAtomValue(this.ref(param) as Atom<Value>) as Value }
  setter(param?: any) { return useSetAtom(this.ref(param) as Writable<Value>) }
}
// Expose a creator
export function atomWithWrapper<Value, FamilyParam = any>(store: Store<Value, FamilyParam>) {
  return new State<Value, FamilyParam>(store)
}

• A "Factory" utility to create atoms :

type StateWrite = (get: Getter, set: Setter, payload: any) => any

export class Factory {
  static persistedStore<Value>(key: string, initialValue: Value) {
    return atomWithStorage<Value>(key, initialValue, storage)
  }
  static store<Value>(read: ((get: Getter) => void) | Value, write?: SetStateAction<Value>) {
    if (!write)
      return atom(read) as Atom<Value>
    return atom(read, write) as WritableAtom<Value, SetStateAction<Value>, any>
  }
  static action<Value>(write: StateWrite) {
    return atom(null, write)
  }
  static family<Param, Value>(
    creator: (param: Param) => Atom<Value>,
    compare?: (a: Param, b: Param) => boolean) {
    return atomFamily<Param, Atom<Value>>(creator, compare)
  }
  static wrapped<Value, FamilyParam = any>(store: Store<Value, FamilyParam>) {
    return new State<Value, FamilyParam>(store)
  }
}

We can now produce clean logic in a "module" without all the imports :

import { Factory, State } from './state'

const count = Factory.persistedStore<number>('persist-count', 0)
const increase = Factory.action(async (get, set) => {
  set(count, prev => prev + 1)
})
const countCreator = Factory.family<number, string>(val => Factory.store(val.toString()))
const wrappedCount = Factory.wrapped(count) // or "atomWithWrapper(count)"

// Expose to the real world
export const CounterModule = {
  count: Factory.wrapped(count),
  countCreator: Factory.wrapped(countCreator),
  wrappedCount, // already wrapped,

  increase: new State(increase), // If you don't like the "Factory.wrapped" nor the "atomWithWrapper"
}

And finally the usage in components :

import { CounterModule } from "modules"

const Component = () => {
  const count = CounterModule.count.get()
  const press = CounterModule.increase.setter()
  const memoizedCopy = CounterModule.countCreator.get(count)
  const wrapped = CounterModule.wrappedCount.get()
  return (
    <View>
      <Text>Counter: {count}</Text>
      <Button
        text="Increase me"
        onPress={press}
        />
    </View>
  )
}

Contraints respected. I'm pretty happy with the minimalism of the API. But what are you thoughts about this pattern? Am I not thinking jotai-ish enough?

Anyway, I found that it could actually be an util atomWithWrapper that would wrap your atom in a class to avoid imports, the atom carries his getter, setter & listener all by itself.
I'm still struggling hard with the typescript tho. I don't know how to easily wrap the call of atom without recopying all the same parameters (could be huge and change over time)

What you guys think?

EDIT: Forgot about the atomFamily stuff, maybe it doesn't fit, gonna check
EDIT2: Added atomFamily stuff, still okay. But broke the typing. AtomFamily is not exported from utils grrr
EDIT3: Fixed typings after adding atomFamily
EDIT4: Added example with atomWithWrapper
EDIT5: Total rework after dai-shi advices

Module logic:

import { persist } from './utilis'
import {
  atom, useAtom,
  useAtomValue, useSetAtom,
} from "jotai"

const someUtil = persist('persist-key', 0)
const count = atom(0)

export const CounterModule = {
  useCount: () => useAtomValue(count),
  useSetCount: () => useSetAtom(count),
}

Usage:

import { CounterModule } from "modules"

const Component = () => {
  const count = CounterModule.useCount()
  const setCount = CounterModule.useSetCount()
  const press = () => {
    setCount(prev => prev + 1)
  }
  return (
    <View>
      <Text>{count}</Text>
      <Button
        text="Increase me"
        onPress={press}
        />
    </View>
  )
}

[dai-shi]

I think it's just a preference, but I would do something like this.

type SetStateAction<Value> = Value | ((prev: Value) => Value)

type State<Value> = {
  useValue: () => Value;
  useSetValue: () => (action: SetStateAction<Value>) => void;
};

const createState = <Value>(initialValue: Value): State<Value> => {
  const a = atom(initialValue);
  return {
    useValue: () => useAtomValue(a),
    useSetValue: () => useSetAtom(a),
  };
};

It doesn't fulfill your constrains though.
btw, I wouldn't overuse atomFamily, because unless you use it with a finite number of params, there's a memory leak risk.

[TwistedMinda]

@dai-shi

My perspective is from a library author, which is a little different from app developers

Yep, totally! You're fulfilling your role perfectly! Just wanted to relativize about the "boost" :p

(edit) you edited, so context lost. 😄

haha yes sorry.
About your example, just to be sure I understood correclty, you mean that you may want to call functions conditonally so you need their reference not just their return value ?
Like, briefly:

// Anywhere in code
const x = someCondition ? CounterModule.useDerived : CounterModule.useNormal

const Component = () => {
  // Inside a component
  const value = x()
}

Anyway, feels really clear to me that using "get" keyword on atoms was a bad idea, now :p
Thanks again!

[dai-shi]

I meant it's inside render in my example.

const Component = () => {
  const x = someCondition && CounterModule.derived; // violating rules of hooks
  return <>{x}</>;
};

[TwistedMinda]

@dai-shi Sorry but would it not also violate rules of hooks with

const Component = () => {
  const x = someCondition && CounterModule.useDerived(); // Error: "React rendered fewer hooks than expected" when "someCondition" changes
  return <>{x}</>;
};

[dai-shi]

Yes, but that can be captured by eslint-plugin-react-hooks. that's the whole point.

[TwistedMinda]

@dai-shi Oh, okay. I get it. With the "get" keywoord i'd prevent potential eslint (or other) plugins that enforce the rules of Hooks directly inside your IDE/Text Editor to work correctly.
Very interesting, thanks.