Merge pull request #239 from zh-lx/feat-add-dict

feat: 优化自定义 dict 的逻辑
zh-lx · May 28, 2024 · 6f57f4b · 6f57f4b
2 parents ae73e88 + b6f1010
commit 6f57f4b
Show file tree

Hide file tree

Showing 3 changed files with 116 additions and 10 deletions.
diff --git a/lib/core/dict/index.ts b/lib/core/dict/index.ts
@@ -1,5 +1,9 @@
 import { Priority, Probability } from "@/common/constant";
 import { Pattern, acTree } from "@/common/segmentit";
+import { getStringLength } from "@/common/utils";
+import DICT1 from "@/data/dict1";
+
+const DefaultName = Symbol("default");
 
 type DICT = {
   [key: string]:
@@ -9,34 +13,92 @@ type DICT = {
     | [string, number, string]; // [拼音, 词频概率, 词性]
 };
 
-export function addDict(dict: DICT | {}, name?: string) {
+type DictOptions = {
+  name?: string;
+  dict1?: "add" | "replace" | "ignore";
+};
+
+const originDictMap = new Map<string | Symbol, { [word: string]: string }>();
+
+export function addDict(dict: DICT | {}, options?: string | DictOptions) {
   const patterns: Pattern[] = [];
+  // string 类型时：options 为 name（）
+  const name = typeof options === "object" ? options.name : options;
+  const dictName = name || DefaultName;
+  const dict1Handle = (options as DictOptions)?.dict1 || "add";
   for (let key in dict as DICT) {
     const value = (dict as DICT)[key];
+    const pinyin = Array.isArray(value) ? value[0] : value;
+    if (getStringLength(key) === 1) {
+      addToOriginDict(
+        dictName,
+        key,
+        pinyin,
+        dict1Handle
+      );
+    }
     if (Array.isArray(value)) {
       patterns.push({
         zh: key,
-        pinyin: value[0],
-        probability: typeof value[1] === "number" ? value[1] : Probability.DICT * key.length * key.length,
+        pinyin,
+        probability:
+          typeof value[1] === "number"
+            ? value[1]
+            : Probability.DICT * key.length * key.length,
         length: key.length,
         priority: Priority.Normal,
-        dict: name || Symbol(""),
+        dict: dictName,
         pos: typeof value[2] === "string" ? value[2] : "",
       });
     } else {
       patterns.push({
         zh: key,
-        pinyin: value,
+        pinyin,
         probability: Probability.DICT * key.length * key.length,
         length: key.length,
         priority: Priority.Normal,
-        dict: name || Symbol(""),
+        dict: dictName,
       });
     }
   }
   acTree.build(patterns);
 }
 
-export function removeDict(dictName: string) {
-  acTree.removeDict(dictName);
+export function removeDict(dictName?: string) {
+  acTree.removeDict(dictName || DefaultName);
+  removeOriginDict(dictName || DefaultName);
+}
+
+function addToOriginDict(
+  dict: string | Symbol,
+  key: string,
+  pinyin: string,
+  handle: "add" | "replace" | "ignore" = "add"
+) {
+  if (!originDictMap.get(dict)) {
+    originDictMap.set(dict, {})
+  }
+  const originDict = originDictMap.get(dict)!;
+  const code = key.charCodeAt(0);
+  if (!originDict[key]) {
+    originDict[key] = DICT1[code] as string;
+  }
+  if (handle === "add") {
+    if (DICT1[code] && !DICT1[code].split(' ').includes(pinyin)) {
+      DICT1[code] += ` ${pinyin}`;
+    } else if (!DICT1[code]) {
+      DICT1[code] = pinyin;
+    }
+  } else if (handle === "replace") {
+    DICT1[code] = pinyin;
+  }
+}
+
+function removeOriginDict(dict: string | Symbol) {
+  const originDict = originDictMap.get(dict) || {};
+  for (let key in originDict) {
+    const code = key.charCodeAt(0);
+    DICT1[code] = originDict[key];
+    delete originDict[key];
+  }
 }
diff --git a/test/dict.test.js b/test/dict.test.js
@@ -20,4 +20,44 @@ describe("addDict", () => {
     const result = pinyin("汉语拼音");
     expect(result).to.be.equal("hàn yǔ pīn yīn");
   });
+
+  it("[addDict]dict handle add", () => {
+    const stringDict = {
+      汉: ['yīn']
+    }
+    addDict(stringDict, { name: "handle-add", dict1: 'add' });
+    const result = pinyin("汉", { multiple: true });
+    expect(result).to.be.equal("hàn yīn");
+    removeDict('handle-add');
+  });
+
+  it("[addDict]dict handle replace", () => {
+    const stringDict = {
+      汉: ['yīn']
+    }
+    addDict(stringDict, { name: "handle-replace", dict1: 'replace' });
+    const result = pinyin("汉", { multiple: true });
+    expect(result).to.be.equal("yīn");
+    removeDict('handle-replace');
+  });
+
+  it("[addDict]undefined dict", () => {
+    const stringDict = {
+      䃜: 'yī'
+    }
+    addDict(stringDict, { name: "handle-new" });
+    const result = pinyin("䃜");
+    expect(result).to.be.equal("yī");
+    removeDict('handle-new');
+  });
+
+  it("[addDict]unnamed dict", () => {
+    const stringDict = {
+      䃜: 'yī'
+    }
+    addDict(stringDict);
+    const result = pinyin("䃜");
+    expect(result).to.be.equal("yī");
+    removeDict();
+  });
 });
diff --git a/types/core/dict/index.d.ts b/types/core/dict/index.d.ts
@@ -1,6 +1,10 @@
 type DICT = {
     [key: string]: string | [string] | [string, number] | [string, number, string];
 };
-export declare function addDict(dict: DICT | {}, name?: string): void;
-export declare function removeDict(dictName: string): void;
+type DictOptions = {
+    name?: string;
+    dict1?: "add" | "replace" | "ignore";
+};
+export declare function addDict(dict: DICT | {}, options?: string | DictOptions): void;
+export declare function removeDict(dictName?: string): void;
 export {};