From d1bbc8f24d63a06dd9f563d8bfe551feb3d75749 Mon Sep 17 00:00:00 2001
From: Hans Pagel <hans.pagel@ueber.io>
Date: Fri, 22 Jan 2021 22:37:43 +0100
Subject: [PATCH] add createTable

---
 docs/src/demos/Nodes/Table/index.vue          |  4 +--
 packages/extension-table/src/table.ts         | 36 +++++++------------
 .../src/utilities/createCell.ts               | 13 +++++++
 .../src/utilities/createTable.ts              | 31 ++++++++++++++++
 .../src/utilities/getTableNodeTypes.ts        | 22 ++++++++++++
 5 files changed, 81 insertions(+), 25 deletions(-)
 create mode 100644 packages/extension-table/src/utilities/createCell.ts
 create mode 100644 packages/extension-table/src/utilities/createTable.ts
 create mode 100644 packages/extension-table/src/utilities/getTableNodeTypes.ts

diff --git a/docs/src/demos/Nodes/Table/index.vue b/docs/src/demos/Nodes/Table/index.vue
index 81cd9cbe..5a176c0f 100644
--- a/docs/src/demos/Nodes/Table/index.vue
+++ b/docs/src/demos/Nodes/Table/index.vue
@@ -1,7 +1,7 @@
 <template>
   <div v-if="editor">
-    <button disabled>
-      ⚠️ createTable
+    <button @click="editor.chain().focus().createTable({ rows: 3, cols: 3, withHeaderRow: true }).run()">
+      ✅ createTable
     </button>
     <button @click="editor.chain().focus().addColumnBefore().run()">
       ✅ addColumnBefore
diff --git a/packages/extension-table/src/table.ts b/packages/extension-table/src/table.ts
index 72e79130..c17cf58a 100644
--- a/packages/extension-table/src/table.ts
+++ b/packages/extension-table/src/table.ts
@@ -16,9 +16,12 @@ import {
   toggleHeaderColumn,
   toggleHeaderRow,
   toggleHeaderCell,
-  setCellAttr,
+  // setCellAttr,
   fixTables,
+  CellSelection,
 } from 'prosemirror-tables'
+import { TextSelection } from 'prosemirror-state'
+import { createTable } from './utilities/createTable'
 import { TableView } from './TableView'
 
 export interface TableOptions {
@@ -54,65 +57,52 @@ export const Table = Node.create({
 
   addCommands() {
     return {
-      // createTable: ({ rowsCount, colsCount, withHeaderRow }) => (
-      //   (state, dispatch) => {
-      //     const offset = state.tr.selection.anchor + 1
+      createTable: ({ rows, cols, withHeaderRow }): Command => ({ state, dispatch }) => {
+        const offset = state.tr.selection.anchor + 1
 
-      //     const nodes = createTable(schema, rowsCount, colsCount, withHeaderRow)
-      //     const tr = state.tr.replaceSelectionWith(nodes).scrollIntoView()
-      //     const resolvedPos = tr.doc.resolve(offset)
+        const nodes = createTable(this.editor.schema, rows, cols, withHeaderRow)
+        const tr = state.tr.replaceSelectionWith(nodes).scrollIntoView()
+        const resolvedPos = tr.doc.resolve(offset)
 
-      //     tr.setSelection(TextSelection.near(resolvedPos))
+        tr.setSelection(TextSelection.near(resolvedPos))
 
-      //     dispatch(tr)
-      //   }
-      // ),
+        return dispatch(tr)
+      },
       addColumnBefore: (): Command => ({ state, dispatch }) => {
-        console.log('addColumnBefore')
         return addColumnBefore(state, dispatch)
       },
       addColumnAfter: (): Command => ({ state, dispatch }) => {
-        console.log('addColumnAfter')
         return addColumnAfter(state, dispatch)
       },
       deleteColumn: (): Command => ({ state, dispatch }) => {
-        console.log('deleteColumn')
         return deleteColumn(state, dispatch)
       },
       addRowBefore: (): Command => ({ state, dispatch }) => {
-        console.log('addRowBefore')
         return addRowBefore(state, dispatch)
       },
       addRowAfter: (): Command => ({ state, dispatch }) => {
-        console.log('addRowAfter')
         return addRowAfter(state, dispatch)
       },
       deleteRow: (): Command => ({ state, dispatch }) => {
-        console.log('deleteRow')
         return deleteRow(state, dispatch)
       },
       deleteTable: (): Command => ({ state, dispatch }) => {
-        console.log('deleteTable')
         return deleteTable(state, dispatch)
       },
       mergeCells: (): Command => ({ state, dispatch }) => {
-        console.log('mergeCells')
+        console.log('mergeCells', { state }, state.selection instanceof CellSelection)
         return mergeCells(state, dispatch)
       },
       splitCell: (): Command => ({ state, dispatch }) => {
-        console.log('splitCell')
         return splitCell(state, dispatch)
       },
       toggleHeaderColumn: (): Command => ({ state, dispatch }) => {
-        console.log('toggleHeaderColumn')
         return toggleHeaderColumn(state, dispatch)
       },
       toggleHeaderRow: (): Command => ({ state, dispatch }) => {
-        console.log('toggleHeaderRow')
         return toggleHeaderRow(state, dispatch)
       },
       toggleHeaderCell: (): Command => ({ state, dispatch }) => {
-        console.log('toggleHeaderCell')
         return toggleHeaderCell(state, dispatch)
       },
       fixTables: (): Command => ({ state, dispatch }) => {
diff --git a/packages/extension-table/src/utilities/createCell.ts b/packages/extension-table/src/utilities/createCell.ts
new file mode 100644
index 00000000..5aebc223
--- /dev/null
+++ b/packages/extension-table/src/utilities/createCell.ts
@@ -0,0 +1,13 @@
+import {
+  NodeType, Fragment,
+  Node as ProsemirrorNode,
+  Schema,
+} from 'prosemirror-model'
+
+export function createCell(cellType: NodeType, cellContent?: Fragment<Schema> | ProsemirrorNode<Schema> | Array<ProsemirrorNode<Schema>>) {
+  if (cellContent) {
+    return cellType.createChecked(null, cellContent)
+  }
+
+  return cellType.createAndFill()
+}
diff --git a/packages/extension-table/src/utilities/createTable.ts b/packages/extension-table/src/utilities/createTable.ts
new file mode 100644
index 00000000..18b41887
--- /dev/null
+++ b/packages/extension-table/src/utilities/createTable.ts
@@ -0,0 +1,31 @@
+import {
+  Schema,
+  Fragment,
+  Node as ProsemirrorNode,
+} from 'prosemirror-model'
+import { createCell } from './createCell'
+
+import { getTableNodeTypes } from './getTableNodeTypes'
+
+export function createTable(schema: Schema, rowsCount: 3, colsCount: 3, withHeaderRow: true, cellContent?: Fragment<Schema> | ProsemirrorNode<Schema> | Array<ProsemirrorNode<Schema>>) {
+  const types = getTableNodeTypes(schema)
+
+  const headerCells = []
+  const cells = []
+
+  for (let index = 0; index < colsCount; index += 1) {
+    cells.push(createCell(types.cell, cellContent))
+
+    if (withHeaderRow) {
+      headerCells.push(createCell(types.header_cell, cellContent))
+    }
+  }
+
+  const rows = []
+
+  for (let index = 0; index < rowsCount; index += 1) {
+    rows.push(types.row.createChecked(null, withHeaderRow && index === 0 ? headerCells : cells))
+  }
+
+  return types.table.createChecked(null, rows)
+}
diff --git a/packages/extension-table/src/utilities/getTableNodeTypes.ts b/packages/extension-table/src/utilities/getTableNodeTypes.ts
new file mode 100644
index 00000000..e0f7324b
--- /dev/null
+++ b/packages/extension-table/src/utilities/getTableNodeTypes.ts
@@ -0,0 +1,22 @@
+import { Schema } from 'prosemirror-model'
+
+export function getTableNodeTypes(schema: Schema) {
+  if (schema.cached.tableNodeTypes) {
+    return schema.cached.tableNodeTypes
+  }
+
+  const roles = {}
+
+  Object.keys(schema.nodes).forEach(type => {
+    const nodeType = schema.nodes[type]
+
+    if (nodeType.spec.tableRole) {
+      // @ts-ignore
+      roles[nodeType.spec.tableRole] = nodeType
+    }
+  })
+
+  schema.cached.tableNodeTypes = roles
+
+  return roles
+}