From c1110efa1adfea318bbb055a5e2c22cfbdae9d98 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@otter.ai>
Date: Tue, 30 Jul 2024 06:23:59 -0700
Subject: [PATCH] feat(atomFamily): supports instanceof

---
 src/vanilla/utils/atomFamily.ts        |  7 ++-
 tests/vanilla/utils/atomFamily.test.ts | 69 ++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 1 deletion(-)
 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..ad8b3720c1 100644
--- a/src/vanilla/utils/atomFamily.ts
+++ b/src/vanilla/utils/atomFamily.ts
@@ -42,11 +42,16 @@ export function atomFamily<Param, AtomType extends Atom<unknown>>(
       }
     }
 
-    const newAtom = initializeAtom(param)
+    const newAtom = Object.setPrototypeOf(
+      initializeAtom(param),
+      createAtom.prototype,
+    )
     atoms.set(param, [newAtom, Date.now()])
     return newAtom
   }
 
+  createAtom.prototype = {}
+
   createAtom.remove = (param: Param) => {
     if (areEqual === undefined) {
       atoms.delete(param)
diff --git a/tests/vanilla/utils/atomFamily.test.ts b/tests/vanilla/utils/atomFamily.test.ts
new file mode 100644
index 0000000000..e6c079ac8f
--- /dev/null
+++ b/tests/vanilla/utils/atomFamily.test.ts
@@ -0,0 +1,69 @@
+import { expect, it, vi } from 'vitest'
+import { 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<number, Atom<number>>(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 support instanceof', () => {
+  const aFamilyA = atomFamily((param: number) => atom(param))
+  const aFamilyB = atomFamily((param: number) => atom(param))
+  expect(aFamilyA(1) instanceof aFamilyA).toEqual(true)
+  expect(aFamilyA(2) instanceof aFamilyA).toEqual(true)
+  expect(aFamilyA(1) instanceof aFamilyB).toEqual(false)
+  expect(aFamilyA(1) instanceof atom).toEqual(false)
+  expect(atom(2) instanceof aFamilyA).toEqual(false)
+})