From e9602626b73b99d6dd57060e9dbd391168ddbabd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philipp=20Ku=CC=88hn?= <kontakt@philipp-kuehn.com>
Date: Wed, 18 Nov 2020 15:18:30 +0100
Subject: [PATCH] improve resetNodeAttributes command

---
 docs/src/demos/Extensions/TextAlign/index.vue |  2 +-
 .../core/src/commands/removeNodeAttributes.ts | 18 +++++++++++++++++
 .../core/src/commands/resetNodeAttributes.ts  | 20 +++++++------------
 packages/core/src/extensions/commands.ts      |  2 +-
 packages/core/src/utils/deleteProps.ts        | 20 +++++++++++++++++++
 packages/extension-text-align/src/index.ts    |  2 +-
 6 files changed, 48 insertions(+), 16 deletions(-)
 create mode 100644 packages/core/src/commands/removeNodeAttributes.ts
 create mode 100644 packages/core/src/utils/deleteProps.ts

diff --git a/docs/src/demos/Extensions/TextAlign/index.vue b/docs/src/demos/Extensions/TextAlign/index.vue
index 53564e23..0ecb20be 100644
--- a/docs/src/demos/Extensions/TextAlign/index.vue
+++ b/docs/src/demos/Extensions/TextAlign/index.vue
@@ -12,7 +12,7 @@
     <button @click="editor.chain().focus().setTextAlign('justify').run()">
       justify
     </button>
-    <button @click="editor.chain().focus().resetNodeAttributes(['textAlign']).run()">
+    <button @click="editor.chain().focus().unsetTextAlign().run()">
       set default
     </button>
     <editor-content :editor="editor" />
diff --git a/packages/core/src/commands/removeNodeAttributes.ts b/packages/core/src/commands/removeNodeAttributes.ts
new file mode 100644
index 00000000..6c10d9af
--- /dev/null
+++ b/packages/core/src/commands/removeNodeAttributes.ts
@@ -0,0 +1,18 @@
+import { NodeType } from 'prosemirror-model'
+import getNodeType from '../utils/getNodeType'
+import deleteProps from '../utils/deleteProps'
+import { Command } from '../types'
+
+export default (typeOrName: string | NodeType, attributes: string[]): Command => ({ tr, state, dispatch }) => {
+  const type = getNodeType(typeOrName, state.schema)
+  const { selection } = tr
+  const { from, to } = selection
+
+  state.doc.nodesBetween(from, to, (node, pos) => {
+    if (node.type === type && dispatch) {
+      tr.setNodeMarkup(pos, undefined, deleteProps(node.attrs, attributes))
+    }
+  })
+
+  return true
+}
diff --git a/packages/core/src/commands/resetNodeAttributes.ts b/packages/core/src/commands/resetNodeAttributes.ts
index db580917..00a02096 100644
--- a/packages/core/src/commands/resetNodeAttributes.ts
+++ b/packages/core/src/commands/resetNodeAttributes.ts
@@ -1,22 +1,16 @@
+import { NodeType } from 'prosemirror-model'
+import getNodeType from '../utils/getNodeType'
+import deleteProps from '../utils/deleteProps'
 import { Command } from '../types'
 
-export default (attributeNames: string[] = []): Command => ({ tr, state, dispatch }) => {
+export default (typeOrName: string | NodeType, attributes: string | string[]): Command => ({ tr, state, dispatch }) => {
+  const type = getNodeType(typeOrName, state.schema)
   const { selection } = tr
   const { from, to } = selection
 
   state.doc.nodesBetween(from, to, (node, pos) => {
-    if (!node.type.isText) {
-      attributeNames.forEach(name => {
-        const attribute = node.type.spec.attrs?.[name]
-        const defaultValue = attribute?.default
-
-        if (attribute && defaultValue !== undefined && dispatch) {
-          tr.setNodeMarkup(pos, undefined, {
-            ...node.attrs,
-            [name]: defaultValue,
-          })
-        }
-      })
+    if (node.type === type && dispatch) {
+      tr.setNodeMarkup(pos, undefined, deleteProps(node.attrs, attributes))
     }
   })
 
diff --git a/packages/core/src/extensions/commands.ts b/packages/core/src/extensions/commands.ts
index 9d647dc2..0b7d3058 100644
--- a/packages/core/src/extensions/commands.ts
+++ b/packages/core/src/extensions/commands.ts
@@ -91,7 +91,7 @@ export const Commands = Extension.create({
        */
       removeMarks,
       /**
-       * Resets all node attributes to the default value.
+       * Resets node attributes to the default value.
        */
       resetNodeAttributes,
       /**
diff --git a/packages/core/src/utils/deleteProps.ts b/packages/core/src/utils/deleteProps.ts
new file mode 100644
index 00000000..9085dcf0
--- /dev/null
+++ b/packages/core/src/utils/deleteProps.ts
@@ -0,0 +1,20 @@
+/**
+ * Remove a property or an array of properties from an object
+ * @param obj Object
+ * @param key Key to remove
+ */
+export default function deleteProps(obj: { [key: string ]: any }, propOrProps: string | string[]) {
+  const props = typeof propOrProps === 'string'
+    ? [propOrProps]
+    : propOrProps
+
+  return Object
+    .keys(obj)
+    .reduce((newObj: { [key: string ]: any }, prop) => {
+      if (!props.includes(prop)) {
+        newObj[prop] = obj[prop]
+      }
+
+      return newObj
+    }, {})
+}
diff --git a/packages/extension-text-align/src/index.ts b/packages/extension-text-align/src/index.ts
index 7477e145..6568a9bf 100644
--- a/packages/extension-text-align/src/index.ts
+++ b/packages/extension-text-align/src/index.ts
@@ -48,7 +48,7 @@ const TextAlign = Extension.create({
        * Unset the text align attribute
        */
       unsetTextAlign: (): Command => ({ commands }) => {
-        return this.options.types.every(type => commands.updateNodeAttributes(type, { textAlign: null }))
+        return this.options.types.every(type => commands.resetNodeAttributes(type, 'textAlign'))
       },
     }
   },