remark-heading.ts90 lines · main
1import Slugger from 'github-slugger'
2import type { Heading, Root } from 'mdast'
3import type { Data, Transformer } from 'unified'
4import { visit } from 'unist-util-visit'
5import type { VFile } from 'vfile'
6
7import type { TOCItemType } from '../types'
8import { flattenNode } from './remark-utils'
9
10const slugger = new Slugger()
11
12declare module 'mdast' {
13 export interface HeadingData extends Data {
14 hProperties?: {
15 id?: string
16 }
17 }
18}
19
20const regex = /\s*\[#([^]+?)]\s*$/
21
22export interface RemarkHeadingOptions {
23 slug?: (root: Root, heading: Heading, text: string) => string
24
25 /**
26 * Allow custom headings ids
27 *
28 * @defaultValue true
29 */
30 customId?: boolean
31
32 /**
33 * Attach an array of `TOCItemType` to `file.data.toc`
34 *
35 * @defaultValue true
36 */
37 generateToc?: boolean
38}
39
40/**
41 * Add heading ids and extract TOC
42 */
43export function remarkHeading({
44 slug: defaultSlug,
45 customId = true,
46 generateToc = true,
47}: RemarkHeadingOptions = {}): Transformer<Root, Root> {
48 return (root: Root, file: VFile) => {
49 const toc: TOCItemType[] = []
50 slugger.reset()
51
52 visit(root, 'heading', (heading: any) => {
53 heading.data ||= {}
54 heading.data.hProperties ||= {}
55
56 let id = heading.data.hProperties.id
57 const lastNode = heading.children.at(-1)
58
59 if (!id && lastNode?.type === 'text' && customId) {
60 const match = regex.exec(lastNode.value)
61
62 if (match?.[1]) {
63 id = match[1]
64 lastNode.value = lastNode.value.slice(0, match.index)
65 }
66 }
67
68 let flattened: string | undefined
69 if (!id) {
70 flattened ??= flattenNode(heading)
71
72 id = defaultSlug ? defaultSlug(root, heading, flattened) : slugger.slug(flattened)
73 }
74
75 heading.data.hProperties.id = id
76
77 if (generateToc) {
78 toc.push({
79 title: flattened ?? flattenNode(heading),
80 url: `#${id}`,
81 depth: heading.depth,
82 })
83 }
84
85 return 'skip'
86 })
87
88 if (generateToc) file.data.toc = toc
89 }
90}