fix: types for ThemedToken

toeverything · Mar 3, 2024 · 420c424 · 420c424
1 parent b311c77
commit 420c424
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 13 deletions.
diff --git a/packages/blocks/src/_common/adapters/html.ts b/packages/blocks/src/_common/adapters/html.ts
@@ -26,9 +26,11 @@ import {
 import type { ElementContent, Root, Text } from 'hast';
 import rehypeParse from 'rehype-parse';
 import rehypeStringify from 'rehype-stringify';
-import { type IThemedToken, type Lang } from 'shiki';
+import { type BundledLanguage, type ThemedToken } from 'shiki';
 import { unified } from 'unified';
 
+import { isPlaintext } from '../../code-block/utils/code-languages.js';
+import { DARK_THEME, LIGHT_THEME } from '../../code-block/utils/consts.js';
 import { getHighLighter } from '../../code-block/utils/high-lighter.js';
 import {
   highlightCache,
@@ -1057,34 +1059,43 @@ export class HtmlAdapter extends BaseAdapter<Html> {
     }, [] as DeltaInsert<object>[]);
     assertEquals(deltas.length, 1, 'Delta length should be 1 in code block');
     const delta = deltas[0];
-    if (rawLang === 'Plain Text' || rawLang === 'Text' || !rawLang) {
+    if (
+      !rawLang ||
+      typeof rawLang !== 'string' ||
+      isPlaintext(rawLang) ||
+      // The rawLang should not be 'Text' here
+      rawLang === 'Text'
+    ) {
       return [
         {
           type: 'text',
           value: delta.insert,
         } as Text,
       ];
     }
-    const lang = rawLang as Lang;
+    const lang = rawLang as BundledLanguage;
+
     const highlighter = await getHighLighter({
       langs: [lang],
+      themes: [LIGHT_THEME, DARK_THEME],
     });
     const cacheKey: highlightCacheKey = `${delta.insert}-${rawLang}-light`;
     const cache = highlightCache.get(cacheKey);
 
-    let tokens: IThemedToken[];
+    let tokens: Omit<ThemedToken, 'offset'>[];
     if (cache) {
       tokens = cache;
     } else {
-      tokens = highlighter
-        .codeToThemedTokens(delta.insert, lang)
-        .reduce((acc, cur, index) => {
+      tokens = highlighter.codeToTokensBase(delta.insert, { lang }).reduce(
+        (acc, cur, index) => {
           if (index === 0) {
             return cur;
           }
 
           return [...acc, { content: '\n', color: 'inherit' }, ...cur];
-        }, []);
+        },
+        [] as Omit<ThemedToken, 'offset'>[]
+      );
       highlightCache.set(cacheKey, tokens);
     }
 

diff --git a/packages/blocks/src/code-block/affine-code-line.ts b/packages/blocks/src/code-block/affine-code-line.ts
@@ -47,10 +47,9 @@ export class AffineCodeLine extends ShadowlessElement {
     const cacheKey: highlightCacheKey = `${this.delta.insert}-${lang}-${mode}`;
     const cache = highlightCache.get(cacheKey);
 
-    let tokens: ThemedToken[] = [
+    let tokens: Omit<ThemedToken, 'offset'>[] = [
       {
         content: this.delta.insert,
-        offset: 0,
       },
     ];
     if (cache) {

diff --git a/packages/blocks/src/code-block/utils/highlight-cache.ts b/packages/blocks/src/code-block/utils/highlight-cache.ts
@@ -30,6 +30,7 @@ class LRUCache<K, V> {
 
 export type highlightCacheKey = `${string}-${string}-${string}`;
 
-export const highlightCache = new LRUCache<highlightCacheKey, ThemedToken[]>(
-  4000
-);
+export const highlightCache = new LRUCache<
+  highlightCacheKey,
+  Omit<ThemedToken, 'offset'>[]
+>(4000);