feat(TreeData): 新增 traverseNode 遍历节点

fjc0k · Dec 21, 2020 · c5dd7b1 · c5dd7b1
1 parent 8612272
commit c5dd7b1
Show file tree

Hide file tree

Showing 3 changed files with 16,796 additions and 1,999 deletions.
diff --git a/src/utils/TreeData.test.ts b/src/utils/TreeData.test.ts
@@ -203,6 +203,39 @@ describe('TreeData', () => {
     ).toMatchSnapshot()
   })
 
+  test('traverseNode', () => {
+    expect(
+      new TreeData(data)
+        .traverseNode(
+          _ => _.node.id === '2',
+          _ => (_.node.id += '_checked'),
+        )
+        .export(),
+    ).toMatchSnapshot()
+  })
+
+  test('traverseNodeDFS', () => {
+    expect(
+      new TreeData(data)
+        .traverseNodeDFS(
+          _ => _.node.id === '2',
+          _ => (_.node.id += '_checked'),
+        )
+        .export(),
+    ).toMatchSnapshot()
+  })
+
+  test('traverseNodeBFS', () => {
+    expect(
+      new TreeData(data)
+        .traverseNodeBFS(
+          _ => _.node.id === '2',
+          _ => (_.node.id += '_checked'),
+        )
+        .export(),
+    ).toMatchSnapshot()
+  })
+
   // test('大数据', () => {
   //   const data = range(0, 1000).map(i => ({
   //     id: `${i}`,

diff --git a/src/utils/TreeData.ts b/src/utils/TreeData.ts
@@ -208,6 +208,74 @@ export class TreeData<TNode extends TreeDataNode> {
     }
   }
 
+  /**
+   * 遍历节点。
+   *
+   * @param node 节点
+   * @param fn 遍历函数
+   * @param searchStrategy 遍历搜索方式，默认为选项中的遍历搜索方式
+   */
+  traverseNode(
+    node:
+      | OneOrMore<TNode>
+      | ((payload: TreeDataTraverseFnPayload<TNode>) => boolean),
+    fn: OneOrMore<TreeDataTraverseFn<TNode> | false>,
+    searchStrategy: TreeDataSearchStrategy = this.searchStrategy,
+  ) {
+    const nodes: TNode[] =
+      typeof node === 'function'
+        ? this.findNodeAll(node as any)
+        : Array.isArray(node)
+        ? node
+        : [node]
+    let _node: TNode | undefined
+    while ((_node = nodes.pop())) {
+      if (Array.isArray(_node[this.childrenPropName])) {
+        const fns: Array<TreeDataTraverseFn<TNode>> = (Array.isArray(fn)
+          ? fn
+          : [fn]
+        ).filter(fn => typeof fn === 'function') as any
+        for (let i = 0; i < fns.length; i++) {
+          TreeData.traverse<TNode>(
+            _node[this.childrenPropName],
+            this.childrenPropName,
+            searchStrategy,
+            fns[i],
+          )
+        }
+      }
+    }
+    return this
+  }
+
+  /**
+   * 深度优先遍历节点。
+   *
+   * @param fn 遍历函数
+   */
+  traverseNodeDFS(
+    node:
+      | OneOrMore<TNode>
+      | ((payload: TreeDataTraverseFnPayload<TNode>) => boolean),
+    fn: OneOrMore<TreeDataTraverseFn<TNode> | false>,
+  ): this {
+    return this.traverseNode(node, fn, 'DFS')
+  }
+
+  /**
+   * 广度优先遍历节点。
+   *
+   * @param fn 遍历函数
+   */
+  traverseNodeBFS(
+    node:
+      | OneOrMore<TNode>
+      | ((payload: TreeDataTraverseFnPayload<TNode>) => boolean),
+    fn: OneOrMore<TreeDataTraverseFn<TNode> | false>,
+  ): this {
+    return this.traverseNode(node, fn, 'BFS')
+  }
+
   /**
    * 遍历。
    *
@@ -218,18 +286,13 @@ export class TreeData<TNode extends TreeDataNode> {
     fn: OneOrMore<TreeDataTraverseFn<TNode> | false>,
     searchStrategy: TreeDataSearchStrategy = this.searchStrategy,
   ): this {
-    const fns: Array<TreeDataTraverseFn<TNode>> = (Array.isArray(fn)
-      ? fn
-      : [fn]
-    ).filter(fn => typeof fn === 'function') as any
-    for (let i = 0; i < fns.length; i++) {
-      TreeData.traverse<TNode>(
-        this.data,
-        this.childrenPropName,
-        searchStrategy,
-        fns[i],
-      )
-    }
+    this.traverseNode(
+      {
+        [this.childrenPropName]: this.data,
+      } as any,
+      fn,
+      searchStrategy,
+    )
     return this
   }