scripts/chunk.ts
8.5 KB| 1 | import { mkdirSync, readFileSync, writeFileSync } from "fs" |
| 2 | import { dirname, join } from "path" |
| 3 | import MarkdownIt from "markdown-it" |
| 4 | |
| 5 | type BlockKind = |
| 6 | | "heading" |
| 7 | | "thematicBreak" |
| 8 | | "html" |
| 9 | | "code" |
| 10 | | "flow" |
| 11 | |
| 12 | interface Block { |
| 13 | kind: BlockKind |
| 14 | md: string |
| 15 | words: number |
| 16 | } |
| 17 | |
| 18 | interface Chunk { |
| 19 | blocks: Block[] |
| 20 | words: number |
| 21 | } |
| 22 | |
| 23 | export interface ChunkCliOptions { |
| 24 | file: string |
| 25 | maxWords: number |
| 26 | outputDir: string |
| 27 | } |
| 28 | |
| 29 | export interface ChunkResult { |
| 30 | source: string |
| 31 | chunks: number |
| 32 | output_dir: string |
| 33 | frontmatter: boolean |
| 34 | words_per_chunk: number[] |
| 35 | } |
| 36 | |
| 37 | const parser = new MarkdownIt({ html: true }) |
| 38 | |
| 39 | export function formatChunkUsage(command: string): string { |
| 40 | return `Usage: ${command} <file> [--max-words 5000] [--output-dir <dir>]` |
| 41 | } |
| 42 | |
| 43 | export function runChunkCli(args: string[], command = "chunk.ts"): number { |
| 44 | const parsed = parseChunkCliArgs(args) |
| 45 | |
| 46 | if ("help" in parsed) { |
| 47 | console.log(formatChunkUsage(command)) |
| 48 | return 0 |
| 49 | } |
| 50 | |
| 51 | if ("error" in parsed) { |
| 52 | console.error(parsed.error) |
| 53 | console.error(formatChunkUsage(command)) |
| 54 | return 1 |
| 55 | } |
| 56 | |
| 57 | const result = chunkMarkdownFile(parsed.file, { |
| 58 | maxWords: parsed.maxWords, |
| 59 | outputDir: parsed.outputDir, |
| 60 | }) |
| 61 | |
| 62 | console.log(JSON.stringify(result)) |
| 63 | return 0 |
| 64 | } |
| 65 | |
| 66 | export function chunkMarkdownFile( |
| 67 | file: string, |
| 68 | options: { maxWords?: number; outputDir?: string } = {} |
| 69 | ): ChunkResult { |
| 70 | const maxWords = options.maxWords ?? 5000 |
| 71 | const outputDir = options.outputDir ?? "" |
| 72 | |
| 73 | const rawContent = normalizeNewlines(readFileSync(file, "utf-8")) |
| 74 | const { frontmatter, body } = extractFrontmatter(rawContent) |
| 75 | const chunks = buildChunks(parseMarkdown(body), maxWords) |
| 76 | |
| 77 | const dir = outputDir ? join(outputDir, "chunks") : join(dirname(file), "chunks") |
| 78 | mkdirSync(dir, { recursive: true }) |
| 79 | |
| 80 | if (frontmatter) { |
| 81 | writeFileSync(join(dir, "frontmatter.md"), frontmatter) |
| 82 | } |
| 83 | |
| 84 | chunks.forEach((chunk, index) => { |
| 85 | const num = String(index + 1).padStart(2, "0") |
| 86 | writeFileSync(join(dir, `chunk-${num}.md`), chunk.blocks.map(block => block.md).join("\n\n")) |
| 87 | }) |
| 88 | |
| 89 | return { |
| 90 | source: file, |
| 91 | chunks: chunks.length, |
| 92 | output_dir: dir, |
| 93 | frontmatter: Boolean(frontmatter), |
| 94 | words_per_chunk: chunks.map(chunk => chunk.words), |
| 95 | } |
| 96 | } |
| 97 | |
| 98 | function parseChunkCliArgs(args: string[]): |
| 99 | | ChunkCliOptions |
| 100 | | { help: true } |
| 101 | | { error: string } { |
| 102 | let file = "" |
| 103 | let maxWords = 5000 |
| 104 | let outputDir = "" |
| 105 | |
| 106 | for (let index = 0; index < args.length; index += 1) { |
| 107 | const arg = args[index] |
| 108 | |
| 109 | if (arg === "-h" || arg === "--help") { |
| 110 | return { help: true } |
| 111 | } |
| 112 | |
| 113 | if (arg === "--max-words") { |
| 114 | const value = args[index + 1] |
| 115 | if (!value) return { error: "Missing value for --max-words" } |
| 116 | maxWords = parsePositiveInt(value, 0) |
| 117 | if (maxWords <= 0) return { error: `Invalid --max-words value: ${value}` } |
| 118 | index += 1 |
| 119 | continue |
| 120 | } |
| 121 | |
| 122 | if (arg === "--output-dir") { |
| 123 | const value = args[index + 1] |
| 124 | if (!value) return { error: "Missing value for --output-dir" } |
| 125 | outputDir = value |
| 126 | index += 1 |
| 127 | continue |
| 128 | } |
| 129 | |
| 130 | if (arg.startsWith("-")) { |
| 131 | return { error: `Unknown option: ${arg}` } |
| 132 | } |
| 133 | |
| 134 | if (!file) { |
| 135 | file = arg |
| 136 | continue |
| 137 | } |
| 138 | |
| 139 | return { error: `Unexpected positional argument: ${arg}` } |
| 140 | } |
| 141 | |
| 142 | if (!file) { |
| 143 | return { error: "Missing input file" } |
| 144 | } |
| 145 | |
| 146 | return { file, maxWords, outputDir } |
| 147 | } |
| 148 | |
| 149 | function parsePositiveInt(value: string | undefined, fallback: number): number { |
| 150 | if (!value) return fallback |
| 151 | const parsed = Number.parseInt(value, 10) |
| 152 | return Number.isFinite(parsed) && parsed > 0 ? parsed : fallback |
| 153 | } |
| 154 | |
| 155 | function normalizeNewlines(text: string): string { |
| 156 | return text.replace(/^\uFEFF/, "").replace(/\r\n?/g, "\n") |
| 157 | } |
| 158 | |
| 159 | function trimBoundaryBlankLines(text: string): string { |
| 160 | return text.replace(/^\n+/, "").replace(/\n+$/, "") |
| 161 | } |
| 162 | |
| 163 | function extractFrontmatter(content: string): { frontmatter: string; body: string } { |
| 164 | const lines = content.split("\n") |
| 165 | if (lines[0] !== "---") { |
| 166 | return { frontmatter: "", body: content } |
| 167 | } |
| 168 | |
| 169 | for (let index = 1; index < lines.length; index += 1) { |
| 170 | if (lines[index] === "---" || lines[index] === "...") { |
| 171 | return { |
| 172 | frontmatter: lines.slice(0, index + 1).join("\n"), |
| 173 | body: lines.slice(index + 1).join("\n").replace(/^\n+/, ""), |
| 174 | } |
| 175 | } |
| 176 | } |
| 177 | |
| 178 | return { frontmatter: "", body: content } |
| 179 | } |
| 180 | |
| 181 | function parseMarkdown(content: string): Block[] { |
| 182 | if (!content.trim()) return [] |
| 183 | |
| 184 | const lines = content.split("\n") |
| 185 | const tokens = parser.parse(content, {}) |
| 186 | const blocks: Block[] = [] |
| 187 | |
| 188 | for (const token of tokens) { |
| 189 | if (!token.map || token.level !== 0) continue |
| 190 | if (token.nesting !== 1 && token.nesting !== 0) continue |
| 191 | |
| 192 | const [startLine, endLine] = token.map |
| 193 | const md = trimBoundaryBlankLines(lines.slice(startLine, endLine).join("\n")) |
| 194 | if (!md) continue |
| 195 | |
| 196 | blocks.push(makeBlock(tokenTypeToBlockKind(token.type), md)) |
| 197 | } |
| 198 | |
| 199 | if (blocks.length === 0) { |
| 200 | const body = trimBoundaryBlankLines(content) |
| 201 | if (body) { |
| 202 | blocks.push(makeBlock("flow", body)) |
| 203 | } |
| 204 | } |
| 205 | |
| 206 | return blocks |
| 207 | } |
| 208 | |
| 209 | function tokenTypeToBlockKind(tokenType: string): BlockKind { |
| 210 | if (tokenType === "heading_open") return "heading" |
| 211 | if (tokenType === "hr") return "thematicBreak" |
| 212 | if (tokenType === "html_block") return "html" |
| 213 | if (tokenType === "fence" || tokenType === "code_block") return "code" |
| 214 | return "flow" |
| 215 | } |
| 216 | |
| 217 | function makeBlock(kind: BlockKind, md: string): Block { |
| 218 | return { |
| 219 | kind, |
| 220 | md: trimBoundaryBlankLines(md), |
| 221 | words: countWords(md), |
| 222 | } |
| 223 | } |
| 224 | |
| 225 | function buildChunks(blocks: Block[], maxWordsPerChunk: number): Chunk[] { |
| 226 | const sections = splitIntoSections(blocks) |
| 227 | const normalizedBlocks: Block[] = [] |
| 228 | |
| 229 | for (const section of sections) { |
| 230 | const sectionWords = section.reduce((sum, block) => sum + block.words, 0) |
| 231 | if (sectionWords <= maxWordsPerChunk) { |
| 232 | normalizedBlocks.push(makeBlock("flow", section.map(block => block.md).join("\n\n"))) |
| 233 | continue |
| 234 | } |
| 235 | |
| 236 | for (const block of section) { |
| 237 | normalizedBlocks.push(...splitOversizedBlock(block, maxWordsPerChunk)) |
| 238 | } |
| 239 | } |
| 240 | |
| 241 | const chunks: Chunk[] = [] |
| 242 | let currentBlocks: Block[] = [] |
| 243 | let currentWords = 0 |
| 244 | |
| 245 | for (const block of normalizedBlocks) { |
| 246 | if (currentWords + block.words > maxWordsPerChunk && currentBlocks.length > 0) { |
| 247 | chunks.push({ blocks: currentBlocks, words: currentWords }) |
| 248 | currentBlocks = [block] |
| 249 | currentWords = block.words |
| 250 | continue |
| 251 | } |
| 252 | |
| 253 | currentBlocks.push(block) |
| 254 | currentWords += block.words |
| 255 | } |
| 256 | |
| 257 | if (currentBlocks.length > 0) { |
| 258 | chunks.push({ blocks: currentBlocks, words: currentWords }) |
| 259 | } |
| 260 | |
| 261 | return chunks |
| 262 | } |
| 263 | |
| 264 | function splitIntoSections(blocks: Block[]): Block[][] { |
| 265 | const sections: Block[][] = [] |
| 266 | let current: Block[] = [] |
| 267 | |
| 268 | for (const block of blocks) { |
| 269 | if (block.kind === "heading" && current.length > 0) { |
| 270 | sections.push(current) |
| 271 | current = [block] |
| 272 | continue |
| 273 | } |
| 274 | |
| 275 | current.push(block) |
| 276 | } |
| 277 | |
| 278 | if (current.length > 0) { |
| 279 | sections.push(current) |
| 280 | } |
| 281 | |
| 282 | return sections |
| 283 | } |
| 284 | |
| 285 | function splitOversizedBlock(block: Block, maxWordsPerChunk: number): Block[] { |
| 286 | if (block.words <= maxWordsPerChunk) return [block] |
| 287 | |
| 288 | if ( |
| 289 | block.kind === "heading" |
| 290 | || block.kind === "thematicBreak" |
| 291 | || block.kind === "html" |
| 292 | || block.kind === "code" |
| 293 | ) { |
| 294 | return [block] |
| 295 | } |
| 296 | |
| 297 | const lines = block.md.split("\n") |
| 298 | if (lines.length <= 1) { |
| 299 | return [block] |
| 300 | } |
| 301 | |
| 302 | const splitBlocks: Block[] = [] |
| 303 | let buffer: string[] = [] |
| 304 | let bufferWords = 0 |
| 305 | |
| 306 | for (const line of lines) { |
| 307 | const lineWords = countWords(line) |
| 308 | if (bufferWords + lineWords > maxWordsPerChunk && buffer.length > 0) { |
| 309 | splitBlocks.push(makeBlock(block.kind, buffer.join("\n"))) |
| 310 | buffer = [line] |
| 311 | bufferWords = lineWords |
| 312 | continue |
| 313 | } |
| 314 | |
| 315 | buffer.push(line) |
| 316 | bufferWords += lineWords |
| 317 | } |
| 318 | |
| 319 | if (buffer.length > 0) { |
| 320 | splitBlocks.push(makeBlock(block.kind, buffer.join("\n"))) |
| 321 | } |
| 322 | |
| 323 | return splitBlocks |
| 324 | } |
| 325 | |
| 326 | function countWords(text: string): number { |
| 327 | const cleaned = text.replace(/[#*`\[\]()>|_~-]/g, " ") |
| 328 | const cjk = cleaned.match(/[\u4e00-\u9fff\u3400-\u4dbf\uf900-\ufaff]/g) |
| 329 | const latin = cleaned.match(/[a-zA-Z0-9]+/g) |
| 330 | return (cjk?.length ?? 0) + (latin?.length ?? 0) |
| 331 | } |
| 332 | |
| 333 | if (import.meta.main) { |
| 334 | process.exit(runChunkCli(process.argv.slice(2), process.argv[1] ?? "chunk.ts")) |
| 335 | } |
文件内容是平台检查时读出的文本，只供查看。平台不会运行包里的任何脚本或程序。
