Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(core): new dev (rev5) methods #2903

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion src/vanilla.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,11 @@
export { atom } from './vanilla/atom.ts'
export type { Atom, WritableAtom, PrimitiveAtom } from './vanilla/atom.ts'

export { createStore, getDefaultStore } from './vanilla/store.ts'
export {
createStore,
getDefaultStore,
unstable_deriveDevStoreRev5,
} from './vanilla/store.ts'

export type {
Getter,
Expand Down
67 changes: 67 additions & 0 deletions src/vanilla/store.ts
Original file line number Diff line number Diff line change
Expand Up @@ -288,6 +288,11 @@ type DevStoreRev4 = {
dev4_get_mounted_atoms: () => Set<AnyAtom>
dev4_restore_atoms: (values: Iterable<readonly [AnyAtom, AnyValue]>) => void
}
type DevStoreRev5 = {
dev5_get_atom_state: (atom: AnyAtom) => AtomState | undefined
dev5_get_mounted_atoms: () => Set<AnyAtom>
dev5_restore_atoms: (values: Iterable<readonly [AnyAtom, AnyValue]>) => void
}

type Store = {
get: <Value>(atom: Atom<Value>) => Value
Expand All @@ -300,6 +305,7 @@ type Store = {
}

export type INTERNAL_DevStoreRev4 = DevStoreRev4
export type INTERNAL_DevStoreRev5 = DevStoreRev5
export type INTERNAL_PrdStore = Store

const buildStore = (
Expand Down Expand Up @@ -801,6 +807,67 @@ const deriveDevStoreRev4 = (store: Store): Store & DevStoreRev4 => {
return Object.assign(derivedStore, devStore)
}

export function unstable_deriveDevStoreRev5(
store: Store,
): Store & DevStoreRev5 {
const debugMountedAtoms = new Set<AnyAtom>()
let savedGetAtomState: StoreArgs[0]
let inRestoreAtom = 0
const derivedStore = store.unstable_derive(
(getAtomState, atomRead, atomWrite, atomOnMount) => {
savedGetAtomState = getAtomState
return [
(atom) => {
const atomState = getAtomState(atom)
const originalMounted = atomState.h
atomState.h = () => {
originalMounted?.()
if (atomState.m) {
debugMountedAtoms.add(atom)
} else {
debugMountedAtoms.delete(atom)
}
}
return atomState
},
atomRead,
(atom, getter, setter, ...args) => {
if (inRestoreAtom) {
return setter(atom, ...args)
}
return atomWrite(atom, getter, setter, ...args)
},
atomOnMount,
]
},
)
const savedStoreSet = derivedStore.set
const devStore: DevStoreRev5 = {
// store dev methods (these are tentative and subject to change without notice)
dev5_get_atom_state: (atom) => savedGetAtomState(atom),
dev5_get_mounted_atoms: () => debugMountedAtoms,
dev5_restore_atoms: (values) => {
const restoreAtom: WritableAtom<null, [], void> = {
read: () => null,
write: (_get, set) => {
++inRestoreAtom
try {
for (const [atom, value] of values) {
if (hasInitialValue(atom)) {
set(atom as never, value)
}
}
} finally {
--inRestoreAtom
}
},
}
savedStoreSet(restoreAtom)
},
}
return Object.assign(derivedStore, devStore)
}

type PrdOrDevStore = Store | (Store & DevStoreRev4)

export const createStore = (): PrdOrDevStore => {
Expand Down
91 changes: 90 additions & 1 deletion tests/vanilla/storedev.test.tsx
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import { describe, expect, it, vi } from 'vitest'
import { atom, createStore } from 'jotai/vanilla'
import { atom, createStore, unstable_deriveDevStoreRev5 } from 'jotai/vanilla'
import type {
INTERNAL_DevStoreRev4,
INTERNAL_PrdStore,
Expand Down Expand Up @@ -107,3 +107,92 @@ describe('[DEV-ONLY] dev-only methods rev4', () => {
expect(store.get(countAtom)).toBe(1)
})
})

describe('[DEV-ONLY] dev-only methods rev5', () => {
const createDevStore = () => {
const store = createStore()
return unstable_deriveDevStoreRev5(store)
}

it('should get atom value', () => {
const store = createDevStore()
const countAtom = atom(0)
countAtom.debugLabel = 'countAtom'
store.set(countAtom, 1)
expect(store.dev5_get_atom_state(countAtom)?.v).toEqual(1)
})

it('should restore atoms and its dependencies correctly', () => {
const store = createDevStore()
const countAtom = atom(0)
const derivedAtom = atom((get) => get(countAtom) * 2)
store.set(countAtom, 1)
store.dev5_restore_atoms([[countAtom, 2]])
expect(store.get(countAtom)).toBe(2)
expect(store.get?.(derivedAtom)).toBe(4)
})

it('should restore atoms and call store listeners correctly', () => {
const store = createDevStore()
const countAtom = atom(0)
const derivedAtom = atom((get) => get(countAtom) * 2)
const countCb = vi.fn()
const derivedCb = vi.fn()
store.set(countAtom, 2)
const unsubCount = store.sub(countAtom, countCb)
const unsubDerived = store.sub(derivedAtom, derivedCb)
store.dev5_restore_atoms([
[countAtom, 1],
[derivedAtom, 2],
])

expect(countCb).toHaveBeenCalled()
expect(derivedCb).toHaveBeenCalled()
unsubCount()
unsubDerived()
})

it('should return all the mounted atoms correctly', () => {
const store = createDevStore()
const countAtom = atom(0)
countAtom.debugLabel = 'countAtom'
const derivedAtom = atom((get) => get(countAtom) * 2)
const unsub = store.sub(derivedAtom, vi.fn())
store.set(countAtom, 1)
const result = store.dev5_get_mounted_atoms()
expect(
Array.from(result).sort(
(a, b) => Object.keys(a).length - Object.keys(b).length,
),
).toStrictEqual([
{ toString: expect.any(Function), read: expect.any(Function) },
{
toString: expect.any(Function),
init: 0,
read: expect.any(Function),
write: expect.any(Function),
debugLabel: 'countAtom',
},
])
unsub()
})

it("should return all the mounted atoms correctly after they're unsubscribed", () => {
const store = createDevStore()
const countAtom = atom(0)
countAtom.debugLabel = 'countAtom'
const derivedAtom = atom((get) => get(countAtom) * 2)
const unsub = store.sub(derivedAtom, vi.fn())
store.set(countAtom, 1)
unsub()
const result = store.dev5_get_mounted_atoms()
expect(Array.from(result)).toStrictEqual([])
})

it('should restore atoms with custom write function', () => {
const store = createDevStore()
const countAtom = atom(0, () => {})
store.dev5_restore_atoms([[countAtom, 1]])
expect(store.get(countAtom)).toBe(1)
})
})
Loading