From 32b83c36f73cd3176e898288915344064c769e63 Mon Sep 17 00:00:00 2001
From: riexn <eriexn@gmail.com>
Date: Fri, 20 Nov 2020 14:15:09 +0400
Subject: [PATCH] added forEach functionality

---
 README.md            |  6 +++
 src/TreeNode.ts      |  8 ++++
 src/types.ts         |  4 ++
 test/forEach.test.ts | 89 ++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 107 insertions(+)
 create mode 100644 test/forEach.test.ts

diff --git a/README.md b/README.md
index 4845b5d..fcef279 100644
--- a/README.md
+++ b/README.md
@@ -103,6 +103,12 @@ tree.moveToSibling({
 tree.flatten();
 ```
 
+### Loop across every node in the tree
+
+```js
+tree.forEach((node) => {console.log(node)})
+```
+
 # Credits
 
 Heavily inspired by [tree-model](https://www.npmjs.com/package/tree-model)
diff --git a/src/TreeNode.ts b/src/TreeNode.ts
index 94aea0f..7db3228 100644
--- a/src/TreeNode.ts
+++ b/src/TreeNode.ts
@@ -3,6 +3,7 @@ import {
   TreeModel,
   MoveToSiblingProps,
   MoveUnderParentProps,
+  ForEachFunction,
 } from './types';
 import { clamp } from './utils';
 import treeHandler from './treeHandler';
@@ -203,4 +204,11 @@ export class TreeNode<T extends TreeModel> {
     }
     return this.parent.children.indexOf(this);
   }
+
+  public forEach(func: ForEachFunction<T>): void {
+    func(this);
+    this.children.map((child) => {
+      child.forEach(func);
+    });
+  }
 }
diff --git a/src/types.ts b/src/types.ts
index a5d9924..4239fc0 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -13,6 +13,10 @@ export interface PredicateFunction<T extends TreeModel> {
   (node: TreeNode<T>): boolean;
 }
 
+export interface ForEachFunction<T extends TreeModel> {
+  (node: TreeNode<T>): void;
+}
+
 export interface MoveUnderParentProps<T extends TreeModel> {
   node: PredicateFunction<T>;
   toParent: PredicateFunction<T>;
diff --git a/test/forEach.test.ts b/test/forEach.test.ts
new file mode 100644
index 0000000..742c48c
--- /dev/null
+++ b/test/forEach.test.ts
@@ -0,0 +1,89 @@
+import treeHandler from '../src/index';
+
+describe('forEach nodes traverser', () => {
+  test('loop across all nodes, and modify their properties', () => {
+    const dataTree: any = {
+      id: 'A',
+      children: [
+        { id: 'B', children: [] },
+        {
+          id: 'C',
+          children: [{ id: 'K', children: [{ id: 'J', children: [] }] }],
+        },
+        ,
+      ],
+    };
+
+    const dataTreeLooped: any = {
+      id: 'A',
+      title: 'ABC',
+      children: [
+        { id: 'B', title: 'ABC', children: [] },
+        {
+          id: 'C',
+          title: 'ABC',
+          children: [
+            {
+              id: 'K',
+              title: 'ABC',
+              children: [{ id: 'J', title: 'ABC', children: [] }],
+            },
+          ],
+        },
+        ,
+      ],
+    };
+
+    const tree = treeHandler.parse(dataTree);
+    tree.forEach((node) => {
+      node.model = { ...node.model, title: 'ABC' };
+    });
+
+    expect(tree.model).toStrictEqual(dataTreeLooped);
+  });
+
+  test('loop across all nodes, and modify their properties based on a condition', () => {
+    const dataTree: any = {
+      id: 'A',
+      children: [
+        { id: 'B', children: [] },
+        {
+          id: 'C',
+          children: [{ id: 'K', children: [{ id: 'J', children: [] }] }],
+        },
+        ,
+      ],
+    };
+
+    const dataTreeLooped: any = {
+      id: 'A',
+      title: 'ABC',
+      children: [
+        { id: 'B', title: 'ABC', children: [] },
+        {
+          id: 'C',
+          title: 'ABC',
+          children: [
+            {
+              id: 'K',
+              title: 'This is K',
+              children: [{ id: 'J', title: 'ABC', children: [] }],
+            },
+          ],
+        },
+        ,
+      ],
+    };
+
+    const tree = treeHandler.parse(dataTree);
+    tree.forEach((node) => {
+      if (node.model.id === 'K') {
+        node.model = { title: 'This is K' };
+      } else {
+        node.model = { title: 'ABC' };
+      }
+    });
+
+    expect(tree.model).toStrictEqual(dataTreeLooped);
+  });
+});