From 85509ff30bf91b3ec9afa6b0f987f45f3f6ebf7f Mon Sep 17 00:00:00 2001
From: Yuriy Demidov <d3m1d0v@yandex-team.ru>
Date: Fri, 1 Mar 2024 15:25:48 +0300
Subject: [PATCH] fix(CodeBlock): fixed parsing of language from token.info and
 added support for different code_block markup (#202)

---
 .../markdown/CodeBlock/CodeBlock.test.ts      | 12 +++---
 .../CodeBlock/CodeBlockSpecs/index.ts         | 38 +++++++++++++++----
 src/extensions/markdown/CodeBlock/const.ts    |  6 ++-
 src/extensions/markdown/CodeBlock/index.ts    |  7 +++-
 4 files changed, 48 insertions(+), 15 deletions(-)

diff --git a/src/extensions/markdown/CodeBlock/CodeBlock.test.ts b/src/extensions/markdown/CodeBlock/CodeBlock.test.ts
index f9964df2..6a64d8af 100644
--- a/src/extensions/markdown/CodeBlock/CodeBlock.test.ts
+++ b/src/extensions/markdown/CodeBlock/CodeBlock.test.ts
@@ -5,8 +5,7 @@ import {createMarkupChecker} from '../../../../tests/sameMarkup';
 import {ExtensionsManager} from '../../../core';
 import {BaseNode, BaseSpecsPreset} from '../../base/specs';
 
-import {CodeBlockSpecs} from './CodeBlockSpecs';
-import {codeBlockLangAttr, codeBlockNodeName} from './const';
+import {CodeBlockNodeAttr, CodeBlockSpecs, codeBlockNodeName} from './CodeBlockSpecs';
 
 const {schema, parser, serializer} = new ExtensionsManager({
     extensions: (builder) => builder.use(BaseSpecsPreset, {}).use(CodeBlockSpecs, {}),
@@ -24,7 +23,7 @@ describe('CodeBlock extension', () => {
     it('should parse a code block', () =>
         same(
             'Some code:\n\n```\nHere it is\n```\n\nPara',
-            doc(p('Some code:'), cb({[codeBlockLangAttr]: ''}, 'Here it is'), p('Para')),
+            doc(p('Some code:'), cb('Here it is'), p('Para')),
         ));
 
     it('parses an intended code block', () =>
@@ -36,14 +35,17 @@ describe('CodeBlock extension', () => {
     it('should parse a fenced code block with info string', () =>
         same(
             'foo\n\n```javascript\n1\n```',
-            doc(p('foo'), cb({[codeBlockLangAttr]: 'javascript'}, '1')),
+            doc(p('foo'), cb({[CodeBlockNodeAttr.Lang]: 'javascript'}, '1')),
         ));
 
     it('should parse a fenced code block with multiple new lines at the end', () =>
-        same('```\nsome code\n\n\n\n```', doc(cb({[codeBlockLangAttr]: ''}, 'some code\n\n\n'))));
+        same('```\nsome code\n\n\n\n```', doc(cb('some code\n\n\n'))));
 
     // TODO: parsed: doc(paragraph("code\nblock"))
     it.skip('should parse html - pre tag', () => {
         parseDOM(schema, '<pre><code>code\nblock</code></pre>', doc(cb('code\nblock')));
     });
+
+    it('should support different markup', () =>
+        same('~~~\n123\n~~~', doc(cb({[CodeBlockNodeAttr.Markup]: '~~~'}, '123'))));
 });
diff --git a/src/extensions/markdown/CodeBlock/CodeBlockSpecs/index.ts b/src/extensions/markdown/CodeBlock/CodeBlockSpecs/index.ts
index f0a7cc64..31fe48bd 100644
--- a/src/extensions/markdown/CodeBlock/CodeBlockSpecs/index.ts
+++ b/src/extensions/markdown/CodeBlock/CodeBlockSpecs/index.ts
@@ -1,8 +1,14 @@
 import type {ExtensionAuto, YENodeSpec} from '../../../../core';
 import {nodeTypeFactory} from '../../../../utils/schema';
 
+export const CodeBlockNodeAttr = {
+    Lang: 'data-language',
+    Markup: 'data-markup',
+} as const;
+
 export const codeBlockNodeName = 'code_block';
-export const codeBlockLangAttr = 'data-language';
+/** @deprecated Use __CodeBlockNodeAttr__ instead */
+export const codeBlockLangAttr = CodeBlockNodeAttr.Lang;
 export const codeBlockType = nodeTypeFactory(codeBlockNodeName);
 
 export type CodeBlockSpecsOptions = {
@@ -13,7 +19,10 @@ export const CodeBlockSpecs: ExtensionAuto<CodeBlockSpecsOptions> = (builder, op
     builder.addNode(codeBlockNodeName, () => ({
         view: opts.nodeview,
         spec: {
-            attrs: {[codeBlockLangAttr]: {default: 'text'}},
+            attrs: {
+                [CodeBlockNodeAttr.Lang]: {default: ''},
+                [CodeBlockNodeAttr.Markup]: {default: '```'},
+            },
             content: 'text*',
             group: 'block',
             code: true,
@@ -25,13 +34,13 @@ export const CodeBlockSpecs: ExtensionAuto<CodeBlockSpecsOptions> = (builder, op
                     tag: 'pre',
                     preserveWhitespace: 'full',
                     getAttrs: (node) => ({
-                        [codeBlockLangAttr]:
-                            (node as Element).getAttribute(codeBlockLangAttr) || '',
+                        [CodeBlockNodeAttr.Lang]:
+                            (node as Element).getAttribute(CodeBlockNodeAttr.Lang) || '',
                     }),
                 },
             ],
             toDOM({attrs}) {
-                return ['pre', attrs[codeBlockLangAttr] ? attrs : {}, ['code', 0]];
+                return ['pre', attrs, ['code', 0]];
             },
         },
         fromYfm: {
@@ -43,11 +52,14 @@ export const CodeBlockSpecs: ExtensionAuto<CodeBlockSpecsOptions> = (builder, op
             },
         },
         toYfm: (state, node) => {
-            state.write('```' + (node.attrs[codeBlockLangAttr] || '') + '\n');
+            const lang: string = node.attrs[CodeBlockNodeAttr.Lang];
+            const markup: string = node.attrs[CodeBlockNodeAttr.Markup];
+
+            state.write(markup + lang + '\n');
             state.text(node.textContent, false);
             // Add a newline to the current content before adding closing marker
             state.write('\n');
-            state.write('```');
+            state.write(markup);
             state.closeBlock(node);
         },
     }));
@@ -60,7 +72,17 @@ export const CodeBlockSpecs: ExtensionAuto<CodeBlockSpecsOptions> = (builder, op
                 name: codeBlockNodeName,
                 type: 'block',
                 noCloseToken: true,
-                getAttrs: (tok) => ({[codeBlockLangAttr]: tok.info || ''}),
+                getAttrs: (tok) => {
+                    const attrs: Record<string, string> = {
+                        [CodeBlockNodeAttr.Markup]: tok.markup,
+                    };
+                    if (tok.info) {
+                        // like in markdown-it
+                        // https://github.com/markdown-it/markdown-it/blob/d07d585b6b15aaee2bc8f7a54b994526dad4dbc5/lib/renderer.mjs#L36-L37
+                        attrs[CodeBlockNodeAttr.Lang] = tok.info.split(/(\s+)/g)[0];
+                    }
+                    return attrs;
+                },
                 prepareContent: removeNewLineAtEnd, // content of fence blocks contains extra \n at the end
             },
         },
diff --git a/src/extensions/markdown/CodeBlock/const.ts b/src/extensions/markdown/CodeBlock/const.ts
index 1db42ad3..9ce37a7d 100644
--- a/src/extensions/markdown/CodeBlock/const.ts
+++ b/src/extensions/markdown/CodeBlock/const.ts
@@ -1,6 +1,10 @@
 import {codeBlockType} from './CodeBlockSpecs';
 
-export {codeBlockNodeName, codeBlockLangAttr} from './CodeBlockSpecs';
+export {
+    codeBlockNodeName,
+    codeBlockLangAttr,
+    CodeBlockNodeAttr as CodeBlockAttr,
+} from './CodeBlockSpecs';
 export const cbAction = 'toCodeBlock';
 /** @deprecated Use `codeBlockType` instead */
 export const cbType = codeBlockType;
diff --git a/src/extensions/markdown/CodeBlock/index.ts b/src/extensions/markdown/CodeBlock/index.ts
index 0f6649f5..cb4c0ea5 100644
--- a/src/extensions/markdown/CodeBlock/index.ts
+++ b/src/extensions/markdown/CodeBlock/index.ts
@@ -13,7 +13,12 @@ import {cbAction, cbType} from './const';
 import {handlePaste} from './handle-paste';
 
 export {resetCodeblock} from './commands';
-export {codeBlockNodeName, codeBlockLangAttr, codeBlockType} from './CodeBlockSpecs';
+export {
+    codeBlockNodeName,
+    CodeBlockNodeAttr,
+    codeBlockLangAttr,
+    codeBlockType,
+} from './CodeBlockSpecs';
 
 export type CodeBlockOptions = CodeBlockSpecsOptions & {
     codeBlockKey?: string | null;