feat: build-in line-numbers module

package.json

@@ -90,11 +90,12 @@
     "@types/node": ">12",
     "@types/prismjs": "^1.16.6",
     "@types/unist": "*",
+    "hastscript": "^7.2.0",
     "prismjs": "^1.24.1",
     "rehype-parse": "^7 || ^ 8",
     "unist-util-is": "^4 || ^5",
-    "unist-util-select": "^4",
-    "unist-util-visit": "^3 || ^4"
+    "unist-util-select": "^4.0.2",
+    "unist-util-visit": "^4.1.1"
   },
   "peerDependencies": {
     "unified": "^10"

src/append-class-name.ts

@@ -0,0 +1,13 @@
+import { Element } from 'hast'
+
+
+export function appendClassName(node: Element, className: string): void {
+  if (!node.properties) node.properties = {}
+  if (!node.properties.className) node.properties.className = []
+
+  if (typeof node.properties.className === 'string') node.properties.className = [node.properties.className]
+  if (typeof node.properties.className === 'number') node.properties.className = [node.properties.className]
+  if (typeof node.properties.className === 'boolean') node.properties.className = []
+
+  node.properties.className.push(className)
+}

src/checker.ts

@@ -0,0 +1,14 @@
+import { Node } from 'unist'
+import { Element } from 'hast'
+import { Text } from 'mdast'
+
+
+export class Checker {
+  public static isElement(node?: Node | null): node is Element {
+    return !!node && node.type === 'element' && 'tagName' in node
+  }
+
+  public static isText(node: Node | null): node is Text {
+    return !!node && node.type === 'text'
+  }
+}

src/get-lang.ts

@@ -0,0 +1,7 @@
+import { Element } from 'hast'
+
+
+export function getLang(node: Element): string {
+  const lang: string = node.properties?.className?.[0] || ''
+  return lang.replace(/^language-/, '')
+}

src/get-line-number.ts

@@ -0,0 +1,4 @@
+export function getLineNumber(str: string): number {
+  const match = str.match(/\n(?!$)/g)
+  return match ? match.length + 1 : 1
+}

src/index.ts

@@ -1,69 +1,27 @@
-import Prism from 'prismjs'
-import parse from 'rehype-parse'
-import unifiedTypes, { unified } from 'unified'
-import * as mdast from 'mdast'
-import * as hast from 'hast'
-import { Node } from 'unist'
+import 'prismjs'
+import unifiedTypes from 'unified'
+import { Element } from 'hast'
 import { visit } from 'unist-util-visit'
-import { select } from 'unist-util-select'
+import { preElementSelector } from './pre-element-selector'
+import { RehypePrismOptions } from './rehype-prism-options'
+import { parseCodeVisitor } from './parse-code-vistor'
+import { Test } from 'unist-util-is'
 
 
-const getLang = (node: hast.Element): string => {
-  const lang: string = node.properties?.className?.[0] || ''
-  return lang.replace(/^language-/, '')
-}
-
-const visitor = (preNode: hast.Element): void => {
-  const codeNode = select('[tagName=code]', preNode) as (hast.Element | null)
-
-  if (!codeNode) return
-  const lang = getLang(codeNode)
-
-  if (!lang || !Prism.languages[lang]) return
-
-  const textNode = select('text', codeNode) as (mdast.Text | null)
-  if (!textNode) return
-
-  const className = `language-${lang}`
-  const html = Prism.highlight(textNode.value, Prism.languages[lang], lang)
-  const tree = unified()
-    .use(parse, { fragment: true })
-    .parse(html) as unknown as hast.Element
-
-  if (!preNode.properties) preNode.properties = {}
-  if (!preNode.properties.className) preNode.properties.className = []
-  if (typeof preNode.properties.className === 'string') preNode.properties.className = [preNode.properties.className]
-  if (typeof preNode.properties.className === 'number') preNode.properties.className = [preNode.properties.className]
-  if (typeof preNode.properties.className === 'boolean') preNode.properties.className = []
-  preNode.properties.className = [...preNode.properties.className, className]
-
-  codeNode.children = tree.children
-}
-
-const selector = (node: hast.Element): boolean => node.tagName === 'pre'
-
-
-export interface RehypePrismOptions {
-  plugins: (
-    'autolinker'| 'autoloader' | 'command-line' | 'copy-to-clipboard' |
-    'custom-class' | 'data-uri-highlight' | 'diff-highlight' |
-    'download-button' | 'file-highlight' | 'filter-highlight-all' |
-    'highlight-keywords' | 'inline-color' | 'jsonp-highlight' |
-    'keep-markup' | 'line-highlight' | 'line-numbers' | 'match-braces' |
-    'normalize-whitespace' | 'previewers' | 'remove-initial-line-feed' |
-    'show-invisibles' | 'show-language' | 'toolbar' | 'treeview' |
-    'unescaped-markup' | 'wpd'
-  )[]
-}
-
-const rehypePrism: unifiedTypes.Plugin<[RehypePrismOptions?]> = (options?: RehypePrismOptions) => {
+const rehypePrism: unifiedTypes.Plugin<[RehypePrismOptions?], Element> = (options?: RehypePrismOptions) => {
   if (options && options.plugins) {
-    for (const plugin of options.plugins) {
+    const plugins = options.plugins.filter(plugin => plugin !== 'line-numbers')
+
+    for (const plugin of plugins) {
       import(`prismjs/plugins/${plugin}/prism-${plugin}.js`)
     }
   }
 
-  return (tree: Node) => visit(tree, selector as any, visitor as any)
+  return tree => visit<Element, Test>(
+    tree,
+    preElementSelector(),
+    parseCodeVisitor(options),
+  )
 }
 
 export default rehypePrism

src/parse-code-vistor.ts

@@ -0,0 +1,57 @@
+import Prism from 'prismjs'
+import rehypeParse from 'rehype-parse'
+import { Element, Parent, ElementContent } from 'hast'
+import { Text } from 'mdast'
+import { Visitor } from 'unist-util-visit/complex-types'
+import { unified } from 'unified'
+import { h } from 'hastscript'
+import { select } from 'unist-util-select'
+import { getLang } from './get-lang'
+import { Checker } from './checker'
+import { appendClassName } from './append-class-name'
+import { getLineNumber } from './get-line-number'
+import { RehypePrismOptions } from './rehype-prism-options'
+
+
+const parser = unified()
+  .use(rehypeParse, { fragment: true })
+
+export function parseCodeVisitor(options?: RehypePrismOptions): Visitor<ElementContent, Parent> {
+  return node => {
+    if (!Checker.isElement(node) || node.tagName !== 'pre') return
+    const preNode = node
+
+    const codeNode = select('[tagName=code]', preNode)
+    if (!Checker.isElement(codeNode)) return
+
+    const lang = getLang(codeNode)
+    if (!lang || !Prism.languages[lang]) return
+
+    const textNode = select('text', codeNode) as (Text | null)
+    if (!Checker.isText(textNode)) return
+
+    const codeText = textNode.value
+
+    const html = Prism.highlight(codeText, Prism.languages[lang], lang)
+    const tree = parser.parse(html) as unknown as Element
+
+    appendClassName(preNode, `language-${lang}`)
+
+    codeNode.children = [...tree.children]
+
+    if (options?.plugins?.includes('line-numbers')) {
+      appendClassName(preNode, 'line-numbers')
+
+      const lineNumber = getLineNumber(codeText)
+      const lineNumberColumn = h(
+        'span',
+        {
+          'aria-hidden': 'true',
+          className: ['line-numbers-rows'],
+        },
+        new Array(lineNumber).fill(h('span')),
+      )
+      codeNode.children.push(lineNumberColumn)
+    }
+  }
+}

src/pre-element-selector.ts

@@ -0,0 +1,7 @@
+import { Test } from 'unist-util-is'
+import { Checker } from './checker'
+
+
+export function preElementSelector(): Test {
+  return node => Checker.isElement(node) && node.tagName === 'pre'
+}

src/rehype-prism-options.ts

@@ -0,0 +1,14 @@
+export type RehypePrismPlugin = (
+  'autolinker'| 'autoloader' | 'command-line' | 'copy-to-clipboard' |
+  'custom-class' | 'data-uri-highlight' | 'diff-highlight' |
+  'download-button' | 'file-highlight' | 'filter-highlight-all' |
+  'highlight-keywords' | 'inline-color' | 'jsonp-highlight' |
+  'keep-markup' | 'line-highlight' | 'line-numbers' | 'match-braces' |
+  'normalize-whitespace' | 'previewers' | 'remove-initial-line-feed' |
+  'show-invisibles' | 'show-language' | 'toolbar' | 'treeview' |
+  'unescaped-markup' | 'wpd'
+)
+
+export interface RehypePrismOptions {
+  plugins?: RehypePrismPlugin[]
+}