feat: limited class support

When a class instance has the new `immerable` symbol in its own properties, its prototype, or on its constructor, the class instance can be drafted.

When drafted, class instances are basically plain objects, just with a different prototype. No magic here.

OTHER CHANGES:

- Allow symbols as property names in drafts

- Add non-enumerable property support to drafts

- Improved error messages

- Fixed issue where ES5 drafts were not being marked as `finalizing` before being shallow cloned, which resulted in unnecessary draft creation

Author: aleclarson

__tests__/base.js

@@ -1,6 +1,6 @@
 "use strict"
-import {Immer, nothing, original, isDraft} from "../src/index"
-import {shallowCopy} from "../src/common"
+import {Immer, nothing, original, isDraft, immerable} from "../src/index"
+import {each, shallowCopy, isEnumerable} from "../src/common"
 import deepFreeze from "deep-freeze"
 import cloneDeep from "lodash.clonedeep"
 import * as lodash from "lodash"
@@ -259,7 +259,9 @@ function runBaseTest(name, useProxies, autoFreeze, useListener) {
                         produce([], d => {
                             d.x = 3
                         })
-                    }).toThrow(/does not support/)
+                    }).toThrow(
+                        "Immer only supports setting array indices and the 'length' property"
+                    )
                 })
 
                 it("throws when a non-numeric property is deleted", () => {
@@ -269,11 +271,86 @@ function runBaseTest(name, useProxies, autoFreeze, useListener) {
                         produce(baseState, d => {
                             delete d.x
                         })
-                    }).toThrow(/does not support/)
+                    }).toThrow("Immer only supports deleting array indices")
                 })
             }
         })
 
+        it("supports `immerable` symbol on constructor", () => {
+            class One {}
+            One[immerable] = true
+            const baseState = new One()
+            const nextState = produce(baseState, draft => {
+                expect(draft).not.toBe(baseState)
+                draft.foo = true
+            })
+            expect(nextState).not.toBe(baseState)
+            expect(nextState.foo).toBeTruthy()
+        })
+
+        it("preserves symbol properties", () => {
+            const test = Symbol("test")
+            const baseState = {[test]: true}
+            const nextState = produce(baseState, s => {
+                expect(s[test]).toBeTruthy()
+                s.foo = true
+            })
+            expect(nextState).toEqual({
+                [test]: true,
+                foo: true
+            })
+        })
+
+        it("preserves non-enumerable properties", () => {
+            const baseState = {}
+            Object.defineProperty(baseState, "foo", {
+                value: true,
+                enumerable: false
+            })
+            const nextState = produce(baseState, s => {
+                expect(s.foo).toBeTruthy()
+                expect(isEnumerable(s, "foo")).toBeFalsy()
+                s.bar = true
+            })
+            expect(nextState.foo).toBeTruthy()
+            expect(isEnumerable(nextState, "foo")).toBeFalsy()
+        })
+
+        it("throws on computed properties", () => {
+            const baseState = {}
+            Object.defineProperty(baseState, "foo", {
+                get: () => {},
+                enumerable: true
+            })
+            expect(() => {
+                produce(baseState, s => {
+                    // Proxies only throw once a change is made.
+                    if (useProxies) {
+                        s.modified = true
+                    }
+                })
+            }).toThrowError("Immer drafts cannot have computed properties")
+        })
+
+        it("allows inherited computed properties", () => {
+            const proto = {}
+            Object.defineProperty(proto, "foo", {
+                get() {
+                    return this.bar
+                },
+                set(val) {
+                    this.bar = val
+                }
+            })
+            const baseState = Object.create(proto)
+            produce(baseState, s => {
+                expect(s.bar).toBeUndefined()
+                s.foo = {}
+                expect(s.bar).toBeDefined()
+                expect(s.foo).toBe(s.bar)
+            })
+        })
+
         it("can rename nested objects (no changes)", () => {
             const nextState = produce({obj: {}}, s => {
                 s.foo = s.obj
@@ -429,7 +506,9 @@ function runBaseTest(name, useProxies, autoFreeze, useListener) {
                             value: 2
                         })
                     })
-                }).toThrow(/does not support/)
+                }).toThrow(
+                    "Object.defineProperty() cannot be used on an Immer draft"
+                )
             })
 
         it("should handle constructor correctly", () => {
@@ -604,7 +683,7 @@ function runBaseTest(name, useProxies, autoFreeze, useListener) {
             it("throws when Object.setPrototypeOf() is used on a draft", () => {
                 produce({}, draft => {
                     expect(() => Object.setPrototypeOf(draft, Array)).toThrow(
-                        /does not support/
+                        "Object.setPrototypeOf() cannot be used on an Immer draft"
                     )
                 })
             })
@@ -905,12 +984,19 @@ function runBaseTest(name, useProxies, autoFreeze, useListener) {
 }
 
 function testObjectTypes(produce) {
+    class Foo {
+        constructor(foo) {
+            this.foo = foo
+            this[immerable] = true
+        }
+    }
     const values = {
         "empty object": {},
         "plain object": {a: 1, b: 2},
         "object (no prototype)": Object.create(null),
         "empty array": [],
-        "plain array": [1, 2]
+        "plain array": [1, 2],
+        "class instance (draftable)": new Foo(1)
     }
     for (const name in values) {
         const value = values[name]
@@ -976,7 +1062,7 @@ function testLiteralTypes(produce) {
         "boxed string": new String(""),
         "boxed boolean": new Boolean(),
         "date object": new Date(),
-        "class instance": new Foo()
+        "class instance (not draftable)": new Foo()
     }
     for (const name in values) {
         describe(name, () => {
@@ -1006,12 +1092,11 @@ function testLiteralTypes(produce) {
 }
 
 function enumerableOnly(x) {
-    const copy = shallowCopy(x)
-    for (const key in copy) {
-        const value = copy[key]
+    const copy = Array.isArray(x) ? x.slice() : Object.assign({}, x)
+    each(copy, (prop, value) => {
         if (value && typeof value === "object") {
-            copy[key] = enumerableOnly(value)
+            copy[prop] = enumerableOnly(value)
         }
-    }
+    })
     return copy
 }

src/common.js

@@ -3,6 +3,11 @@ export const NOTHING =
         ? Symbol("immer-nothing")
         : {["immer-nothing"]: true}
 
+export const DRAFTABLE =
+    typeof Symbol !== "undefined"
+        ? Symbol("immer-draftable")
+        : "__$immer_draftable"
+
 export const DRAFT_STATE =
     typeof Symbol !== "undefined" ? Symbol("immer-state") : "__$immer_state"
 
@@ -11,11 +16,11 @@ export function isDraft(value) {
 }
 
 export function isDraftable(value) {
-    if (!value) return false
-    if (typeof value !== "object") return false
+    if (!value || typeof value !== "object") return false
     if (Array.isArray(value)) return true
     const proto = Object.getPrototypeOf(value)
-    return proto === null || proto === Object.prototype
+    if (!proto || proto === Object.prototype) return true
+    return !!value[DRAFTABLE] || !!value.constructor[DRAFTABLE]
 }
 
 export function original(value) {
@@ -36,10 +41,39 @@ export const assign =
         return target
     }
 
-export function shallowCopy(value) {
-    if (Array.isArray(value)) return value.slice()
-    const target = value.__proto__ === undefined ? Object.create(null) : {}
-    return assign(target, value)
+export const ownKeys =
+    typeof Reflect !== "undefined"
+        ? Reflect.ownKeys
+        : obj =>
+              Object.getOwnPropertyNames(obj).concat(
+                  Object.getOwnPropertySymbols(obj)
+              )
+
+export function shallowCopy(base, invokeGetters = false) {
+    if (Array.isArray(base)) return base.slice()
+    const clone = Object.create(Object.getPrototypeOf(base))
+    ownKeys(base).forEach(key => {
+        if (key === DRAFT_STATE) {
+            return // Never copy over draft state.
+        }
+        const desc = Object.getOwnPropertyDescriptor(base, key)
+        if (desc.get) {
+            if (!invokeGetters) {
+                throw new Error("Immer drafts cannot have computed properties")
+            }
+            desc.value = desc.get.call(base)
+        }
+        if (desc.enumerable) {
+            clone[key] = desc.value
+        } else {
+            Object.defineProperty(clone, key, {
+                value: desc.value,
+                writable: true,
+                configurable: true
+            })
+        }
+    })
+    return clone
 }
 
 export function each(value, cb) {
@@ -50,6 +84,18 @@ export function each(value, cb) {
     }
 }
 
+export function eachOwn(value, cb) {
+    if (Array.isArray(value)) {
+        for (let i = 0; i < value.length; i++) cb(i, value[i], value)
+    } else {
+        ownKeys(value).forEach(key => cb(key, value[key], value))
+    }
+}
+
+export function isEnumerable(base, prop) {
+    return Object.getOwnPropertyDescriptor(base, prop).enumerable
+}
+
 export function has(thing, prop) {
     return Object.prototype.hasOwnProperty.call(thing, prop)
 }

src/es5.js

@@ -8,7 +8,9 @@ import {
     isDraft,
     isDraftable,
     shallowCopy,
-    DRAFT_STATE
+    DRAFT_STATE,
+    eachOwn,
+    isEnumerable
 } from "./common"
 
 const descriptors = {}
@@ -33,13 +35,16 @@ export function createDraft(base, parent) {
         const state = base[DRAFT_STATE]
         // Avoid creating new drafts when copying.
         state.finalizing = true
-        draft = shallowCopy(state.draft)
+        draft = shallowCopy(state.draft, true)
         state.finalizing = false
     } else {
         draft = shallowCopy(base)
     }
-    each(base, prop => {
-        Object.defineProperty(draft, "" + prop, createPropertyProxy("" + prop))
+
+    const isArray = Array.isArray(base)
+    eachOwn(draft, prop => {
+        const enumerable = isArray || isEnumerable(base, prop)
+        proxyProperty(draft, prop, enumerable)
     })
 
     // See "proxy.js" for property documentation.
@@ -103,20 +108,23 @@ function prepareCopy(state) {
     if (!state.copy) state.copy = shallowCopy(state.base)
 }
 
-function createPropertyProxy(prop) {
-    return (
-        descriptors[prop] ||
-        (descriptors[prop] = {
+function proxyProperty(draft, prop, enumerable) {
+    let desc = descriptors[prop]
+    if (desc) {
+        desc.enumerable = enumerable
+    } else {
+        descriptors[prop] = desc = {
             configurable: true,
-            enumerable: true,
+            enumerable,
             get() {
                 return get(this[DRAFT_STATE], prop)
             },
             set(value) {
                 set(this[DRAFT_STATE], prop, value)
             }
-        })
-    )
+        }
+    }
+    Object.defineProperty(draft, prop, desc)
 }
 
 function assertUnrevoked(state) {

src/immer.d.ts

@@ -122,6 +122,16 @@ declare class Nothing {
  */
 export const nothing: Nothing
 
+/**
+ * To let Immer treat your class instances as plain immutable objects
+ * (albeit with a custom prototype), you must define either an instance property
+ * or a static property on each of your custom classes.
+ *
+ * Otherwise, your class instance will never be drafted, which means it won't be
+ * safe to mutate in a produce callback.
+ */
+export const immerable: unique symbol
+
 /**
  * Pass true to automatically freeze all copies created by Immer.
  *