From cdb5f2567c92ffb494c70f90a47fcb235688f7f9 Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Wed, 31 Jul 2024 00:35:13 -0700 Subject: [PATCH 1/5] feat(atomFamily): support getParams and unstable_listen api --- src/vanilla/utils/atomFamily.ts | 43 +++++++++++- tests/vanilla/utils/atomFamily.test.ts | 94 ++++++++++++++++++++++++++ 2 files changed, 134 insertions(+), 3 deletions(-) create mode 100644 tests/vanilla/utils/atomFamily.test.ts diff --git a/src/vanilla/utils/atomFamily.ts b/src/vanilla/utils/atomFamily.ts index 81b0098613..d97ec41ea8 100644 --- a/src/vanilla/utils/atomFamily.ts +++ b/src/vanilla/utils/atomFamily.ts @@ -1,11 +1,22 @@ import type { Atom } from '../../vanilla.ts' type ShouldRemove = (createdAt: number, param: Param) => boolean +type Cleanup = () => void +type Callback = (event: { + type: 'CREATE' | 'REMOVE' + param: Param + atom: AtomType +}) => void export interface AtomFamily { (param: Param): AtomType + getParams(): Iterable remove(param: Param): void setShouldRemove(shouldRemove: ShouldRemove | null): void + /** + * fires when a atom is created or removed + */ + unstable_listen(callback: Callback): Cleanup } export function atomFamily>( @@ -20,6 +31,7 @@ export function atomFamily>( type CreatedAt = number // in milliseconds let shouldRemove: ShouldRemove | null = null const atoms: Map = new Map() + const listeners = new Set>() const createAtom = (param: Param) => { let item: [AtomType, CreatedAt] | undefined if (areEqual === undefined) { @@ -43,16 +55,40 @@ export function atomFamily>( } const newAtom = initializeAtom(param) + notifyListeners('CREATE', param, newAtom) atoms.set(param, [newAtom, Date.now()]) return newAtom } + function notifyListeners( + type: 'CREATE' | 'REMOVE', + param: Param, + atom: AtomType, + ) { + for (const listener of listeners) { + listener({ type, param, atom }) + } + } + + createAtom.unstable_listen = (callback: Callback) => { + listeners.add(callback) + return () => { + listeners.delete(callback) + } + } + + createAtom.getParams = () => atoms.keys() + createAtom.remove = (param: Param) => { if (areEqual === undefined) { + if (!atoms.has(param)) return + const [atom] = atoms.get(param)! + notifyListeners('REMOVE', param, atom) atoms.delete(param) } else { - for (const [key] of atoms) { + for (const [key, [atom]] of atoms) { if (areEqual(key, param)) { + notifyListeners('REMOVE', key, atom) atoms.delete(key) break } @@ -63,8 +99,9 @@ export function atomFamily>( createAtom.setShouldRemove = (fn: ShouldRemove | null) => { shouldRemove = fn if (!shouldRemove) return - for (const [key, value] of atoms) { - if (shouldRemove(value[1], key)) { + for (const [key, [atom, createdAt]] of atoms) { + if (shouldRemove(createdAt, key)) { + notifyListeners('REMOVE', key, atom) atoms.delete(key) } } diff --git a/tests/vanilla/utils/atomFamily.test.ts b/tests/vanilla/utils/atomFamily.test.ts new file mode 100644 index 0000000000..01d9caf322 --- /dev/null +++ b/tests/vanilla/utils/atomFamily.test.ts @@ -0,0 +1,94 @@ +import { expect, it, vi } from 'vitest' +import { type Atom, atom, createStore } from 'jotai/vanilla' +import { atomFamily } from 'jotai/vanilla/utils' + +it('should create atoms with different params', () => { + const store = createStore() + const aFamily = atomFamily((param: number) => atom(param)) + + expect(store.get(aFamily(1))).toEqual(1) + expect(store.get(aFamily(2))).toEqual(2) +}) + +it('should remove atoms', () => { + const store = createStore() + const initializeAtom = vi.fn((param: number) => atom(param)) + const aFamily = atomFamily(initializeAtom) + + expect(store.get(aFamily(1))).toEqual(1) + expect(store.get(aFamily(2))).toEqual(2) + aFamily.remove(2) + initializeAtom.mockClear() + expect(store.get(aFamily(1))).toEqual(1) + expect(initializeAtom).toHaveBeenCalledTimes(0) + expect(store.get(aFamily(2))).toEqual(2) + expect(initializeAtom).toHaveBeenCalledTimes(1) +}) + +it('should remove atoms with custom comparator', () => { + const store = createStore() + const initializeAtom = vi.fn((param: number) => atom(param)) + const aFamily = atomFamily(initializeAtom, (a, b) => a === b) + + expect(store.get(aFamily(1))).toEqual(1) + expect(store.get(aFamily(2))).toEqual(2) + expect(store.get(aFamily(3))).toEqual(3) + aFamily.remove(2) + initializeAtom.mockClear() + expect(store.get(aFamily(1))).toEqual(1) + expect(initializeAtom).toHaveBeenCalledTimes(0) + expect(store.get(aFamily(2))).toEqual(2) + expect(initializeAtom).toHaveBeenCalledTimes(1) +}) + +it('should remove atoms with custom shouldRemove', () => { + const store = createStore() + const initializeAtom = vi.fn((param: number) => atom(param)) + const aFamily = atomFamily>(initializeAtom) + expect(store.get(aFamily(1))).toEqual(1) + expect(store.get(aFamily(2))).toEqual(2) + expect(store.get(aFamily(3))).toEqual(3) + aFamily.setShouldRemove((_createdAt, param) => param % 2 === 0) + initializeAtom.mockClear() + expect(store.get(aFamily(1))).toEqual(1) + expect(initializeAtom).toHaveBeenCalledTimes(0) + expect(store.get(aFamily(2))).toEqual(2) + expect(initializeAtom).toHaveBeenCalledTimes(1) + expect(store.get(aFamily(3))).toEqual(3) + expect(initializeAtom).toHaveBeenCalledTimes(1) +}) + +it('should notify listeners', () => { + const aFamily = atomFamily((param: number) => atom(param)) + const listener = vi.fn(() => {}) + type Event = { type: 'CREATE' | 'REMOVE'; param: number; atom: Atom } + const unsubscribe = aFamily.unstable_listen(listener) + const atom1 = aFamily(1) + expect(listener).toHaveBeenCalledTimes(1) + const eventCreate = listener.mock.calls[0]?.at(0) as unknown as Event + if (!eventCreate) throw new Error('eventCreate is undefined') + expect(eventCreate.type).toEqual('CREATE') + expect(eventCreate.param).toEqual(1) + expect(eventCreate.atom).toEqual(atom1) + listener.mockClear() + aFamily.remove(1) + expect(listener).toHaveBeenCalledTimes(1) + const eventRemove = listener.mock.calls[0]?.at(0) as unknown as Event + expect(eventRemove.type).toEqual('REMOVE') + expect(eventRemove.param).toEqual(1) + expect(eventRemove.atom).toEqual(atom1) + unsubscribe() + listener.mockClear() + aFamily(2) + expect(listener).toHaveBeenCalledTimes(0) +}) + +it('should return all params', () => { + const store = createStore() + const aFamily = atomFamily((param: number) => atom(param)) + + expect(store.get(aFamily(1))).toEqual(1) + expect(store.get(aFamily(2))).toEqual(2) + expect(store.get(aFamily(3))).toEqual(3) + expect(Array.from(aFamily.getParams())).toEqual([1, 2, 3]) +}) From f7da0cbb326b6810ebe34045ab148bee0d7cb40d Mon Sep 17 00:00:00 2001 From: Daishi Kato Date: Wed, 31 Jul 2024 18:09:13 +0900 Subject: [PATCH 2/5] Update tests/vanilla/utils/atomFamily.test.ts --- tests/vanilla/utils/atomFamily.test.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/vanilla/utils/atomFamily.test.ts b/tests/vanilla/utils/atomFamily.test.ts index 01d9caf322..a66f4d1a7b 100644 --- a/tests/vanilla/utils/atomFamily.test.ts +++ b/tests/vanilla/utils/atomFamily.test.ts @@ -1,5 +1,6 @@ import { expect, it, vi } from 'vitest' -import { type Atom, atom, createStore } from 'jotai/vanilla' +import { atom, createStore } from 'jotai/vanilla' +import type { Atom } from 'jotai/vanilla' import { atomFamily } from 'jotai/vanilla/utils' it('should create atoms with different params', () => { From a1c4da6fae442095148ddc05c125a1337c231ab3 Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Wed, 31 Jul 2024 10:05:47 -0700 Subject: [PATCH 3/5] Update src/vanilla/utils/atomFamily.ts Co-authored-by: Daishi Kato --- src/vanilla/utils/atomFamily.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/src/vanilla/utils/atomFamily.ts b/src/vanilla/utils/atomFamily.ts index d97ec41ea8..36e59ed584 100644 --- a/src/vanilla/utils/atomFamily.ts +++ b/src/vanilla/utils/atomFamily.ts @@ -15,6 +15,7 @@ export interface AtomFamily { setShouldRemove(shouldRemove: ShouldRemove | null): void /** * fires when a atom is created or removed + * This API is for advanced use cases, and can change without notice. */ unstable_listen(callback: Callback): Cleanup } From 4076e74fbde7e704b925a402ee9b7ca6e27df1f4 Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Wed, 31 Jul 2024 10:08:48 -0700 Subject: [PATCH 4/5] call notifyListeners after adding to set --- src/vanilla/utils/atomFamily.ts | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/vanilla/utils/atomFamily.ts b/src/vanilla/utils/atomFamily.ts index 36e59ed584..3ce5bcbb02 100644 --- a/src/vanilla/utils/atomFamily.ts +++ b/src/vanilla/utils/atomFamily.ts @@ -1,4 +1,4 @@ -import type { Atom } from '../../vanilla.ts' +import { type Atom } from '../../vanilla.ts' type ShouldRemove = (createdAt: number, param: Param) => boolean type Cleanup = () => void @@ -56,8 +56,8 @@ export function atomFamily>( } const newAtom = initializeAtom(param) - notifyListeners('CREATE', param, newAtom) atoms.set(param, [newAtom, Date.now()]) + notifyListeners('CREATE', param, newAtom) return newAtom } @@ -84,13 +84,13 @@ export function atomFamily>( if (areEqual === undefined) { if (!atoms.has(param)) return const [atom] = atoms.get(param)! - notifyListeners('REMOVE', param, atom) atoms.delete(param) + notifyListeners('REMOVE', param, atom) } else { for (const [key, [atom]] of atoms) { if (areEqual(key, param)) { - notifyListeners('REMOVE', key, atom) atoms.delete(key) + notifyListeners('REMOVE', key, atom) break } } @@ -102,8 +102,8 @@ export function atomFamily>( if (!shouldRemove) return for (const [key, [atom, createdAt]] of atoms) { if (shouldRemove(createdAt, key)) { - notifyListeners('REMOVE', key, atom) atoms.delete(key) + notifyListeners('REMOVE', key, atom) } } } From 48ed9fad19d8ab37a48eafc07ebe5a508c6dfcb0 Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Wed, 31 Jul 2024 10:11:04 -0700 Subject: [PATCH 5/5] add jsdoc comment for CreatedAt --- src/vanilla/utils/atomFamily.ts | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/vanilla/utils/atomFamily.ts b/src/vanilla/utils/atomFamily.ts index 3ce5bcbb02..4667354890 100644 --- a/src/vanilla/utils/atomFamily.ts +++ b/src/vanilla/utils/atomFamily.ts @@ -1,6 +1,10 @@ import { type Atom } from '../../vanilla.ts' -type ShouldRemove = (createdAt: number, param: Param) => boolean +/** + * in milliseconds + */ +type CreatedAt = number +type ShouldRemove = (createdAt: CreatedAt, param: Param) => boolean type Cleanup = () => void type Callback = (event: { type: 'CREATE' | 'REMOVE' @@ -29,7 +33,6 @@ export function atomFamily>( initializeAtom: (param: Param) => AtomType, areEqual?: (a: Param, b: Param) => boolean, ) { - type CreatedAt = number // in milliseconds let shouldRemove: ShouldRemove | null = null const atoms: Map = new Map() const listeners = new Set>()