atom_family_strict.ts

export function atomFamilyStrict<
  Key,
  Args extends unknown[],
  AtomType extends {
    onMount?: (set: (...args: Args) => void) => void | (() => void);
  },
>(initializeAtom: (param: Key) => AtomType) {
  const atoms = new Map<string, AtomType>();

  const family = (key: Key) => {
    const index = stringify(key);
    const atom = atoms.get(index);
    if (atom) {
      return atom;
    }

    const newAtom = initializeAtom(key);
    if ("onMount" in newAtom) {
      overrideOnMount(newAtom, index);
    }
    atoms.set(index, newAtom);
    return newAtom;
  };
  family.atoms = atoms;
  return family;

  function overrideOnMount(atom: AtomType, key: string) {
    const originalOnMount = atom.onMount;
    atom.onMount = (set: (...args: Args) => void) => {
      atoms.set(key, atom);
      const originalOnUnmount = originalOnMount?.(set);

      return () => {
        atoms.delete(key);
        return originalOnUnmount?.();
      };
    };
  }
}

const ids = new WeakMap<object, number>();
const salt = Math.random().toString(36).slice(2);
let idCount = 0;

export function stringify(key: unknown): string {
  return JSON.stringify(key, (_, val) => {
    if (isPlainObject(val)) {
      return Object.keys(val)
        .sort()
        .reduce((result, key) => (result[key] = val[key], result), {} as Record<string, unknown>);
    }
    return isStringifiable(val)
      ? val
      : `[${salt}:${(ids.get(val) ?? ids.set(val, ++idCount).get(val))}]`;
  });
}

// Copied from: https://github.com/jonschlinkert/is-plain-object
// deno-lint-ignore ban-types
function isPlainObject(o: unknown): o is Object {
  if (!hasObjectPrototype(o)) {
    return false;
  }

  // If has modified constructor
  const ctor = o.constructor;
  if (typeof ctor === "undefined") {
    return true;
  }

  // If has modified prototype
  const proto = ctor.prototype;
  if (!hasObjectPrototype(proto)) {
    return false;
  }

  // If constructor does not have an Object-specific method
  if (!Object.prototype.hasOwnProperty.call(proto, "isPrototypeOf")) {
    return false;
  }

  // Most likely a plain Object
  return true;
}

// deno-lint-ignore ban-types
function hasObjectPrototype(o: unknown): o is Object {
  return Object.prototype.toString.call(o) === "[object Object]";
}

function isStringifiable(o: unknown): boolean {
  const type = typeof o;
  if (type === "object" || type === "function" || type === "symbol") {
    return false;
  }
  return true;
}

atom_with_init.ts

import { atom, type Getter, type Setter } from 'jotai'

type SetAtom<Args extends unknown[], Result> = <A extends Args>(...args: A) => Result
type Read<Value, SetSelf = never> = (
  get: Getter,
  options: {
    readonly signal: AbortSignal
    readonly setSelf: SetSelf
  }
) => Value
type Write<Args extends unknown[], Result> = (get: Getter, set: Setter, ...args: Args) => Result

export function atomWithInit<Value, Args extends unknown[], Result>(
  get: Read<Value, SetAtom<Args, Result>>,
  set: Write<Args, Result>,
  onMount: Write<[], void>
) {
  const initAtom = atom(null, (get, set) => {
    onMount(get, set)
  })
  initAtom.onMount = (set) => set()

  return atom((getter, options) => {
    getter(initAtom)
    get(getter, options)
  }, set)
}

setter_atom.ts

import { atom, WritableAtom } from "jotai";

export function setterAtom<Value, Args extends unknown[], Result>(
  writable: WritableAtom<Value, Args, Result>,
) {
  let resolver: (resolve: (...args: Args) => Result) => void;
  const promise = new Promise<(...args: Args) => Result>((resolve) => resolver = resolve);

  const promiseAtom = atom(promise);
  const writeAtom = atom(
    (get) => get(promiseAtom),
    (_get, set) => resolver((...args: Args) => set(writable, ...args)),
  );
  writeAtom.onMount = (set) => set();
  return writeAtom;
}