fix: match onCommentEnd (#113)

huozhi · Jan 5, 2025 · 7a54df6 · 7a54df6
1 parent 53fd92a
commit 7a54df6
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 7 deletions.
diff --git a/lib/index.d.ts b/lib/index.d.ts
@@ -1,6 +1,7 @@
 type HighlightOptions = {
   keywords?: Set<string>
-  onCommentStart?: (curr: string, next: string) => boolean
+  onCommentStart?: (curr: string, next: string) => number | boolean
+  onCommentEnd?: (curr: string, prev: string) => number | boolean
 }
 
 export function highlight(code: string, options?: HighlightOptions): string

diff --git a/lib/index.js b/lib/index.js
@@ -79,6 +79,7 @@ const Signs = new Set([
 const DefaultOptions = {
   keywords: Keywords_Js,
   onCommentStart: isCommentStart_Js,
+  onCommentEnd: isCommentEnd_Js,
 }
 
 /**
@@ -184,6 +185,10 @@ function isCommentStart_Js(curr, next) {
   return str === '//' || str === '/*'
 }
 
+function isCommentEnd_Js(prev, curr) {
+  return curr === '\n' || (prev + curr) === '*/'
+}
+
 function isRegexStart(str) {
   return str[0] === '/' && !isCommentStart_Js(str[0], str[1])
 }
@@ -197,6 +202,7 @@ function tokenize(code, options) {
   const {
     keywords,
     onCommentStart,
+    onCommentEnd,
   } = { ...DefaultOptions, ...options }
 
   let current = ''
@@ -549,10 +555,13 @@ function tokenize(code, options) {
     } else if (onCommentStart(curr, next)) {
       append()
       const start = i
-      if (next === '/') {
-        for (; i < code.length && code[i] !== '\n'; i++);
-      } else {
-        for (; i < code.length && code[i - 1] + code[i] !== '*/'; i++);
+      const commentType = onCommentStart(curr, next)
+
+      // just match the comment, commentType === true
+      // inline comment, commentType === 1
+      // block comment, commentType === 2
+      if (commentType) {
+        for (; i < code.length && !onCommentEnd(code[i - 1], code[i]); i++);
       }
       current = code.slice(start, i + 1)
       append(T_COMMENT)

diff --git a/test/tokenize.test.ts b/test/tokenize.test.ts
@@ -24,14 +24,26 @@ describe('tokenize - customized keywords', () => {
 
 describe('tokenize - customized comment rule', () => {
   it('should tokenize the input string with the given comment rule', () => {
-    const input = `def f(): return 2 # this is a comment`
+    const input = `\
+    # define a function
+    def f():
+      return 2 # this is a comment
+    `
     const keywords = new Set(['def', 'return'])
     const onCommentStart = (curr, next) => {
       return curr === '#'
     }
-    const actual = getTokensAsString(tokenize(input, { keywords, onCommentStart }))
+    const onCommentEnd = (prev, curr) => {
+      return curr === '\n'
+    }
+    const actual = getTokensAsString(tokenize(input, {
+      keywords, 
+      onCommentStart,
+      onCommentEnd
+    }))
     expect(actual).toMatchInlineSnapshot(`
       [
+        "# define a function => comment",
         "def => keyword",
         "f => identifier",
         "( => sign",