From ae8ab4b56c9963e3749b6ec03e5debb4a210b75c Mon Sep 17 00:00:00 2001
From: baoyang <lbyang222@gmail.com>
Date: Sat, 25 Feb 2017 16:33:22 +0800
Subject: [PATCH] Tree: fix setCheckedKeys() not work on node that has
 childNodes (#2967)

---
 examples/docs/en-US/tree.md           | 29 ++++++------
 examples/docs/zh-CN/tree.md           | 27 ++++++------
 packages/tree/src/model/tree-store.js | 63 +++++++++++++++++++++++----
 test/unit/specs/tree.spec.js          | 56 +++++++++++++++++++++++-
 4 files changed, 136 insertions(+), 39 deletions(-)

diff --git a/examples/docs/en-US/tree.md b/examples/docs/en-US/tree.md
index 7167c0d75..52fc9ad17 100644
--- a/examples/docs/en-US/tree.md
+++ b/examples/docs/en-US/tree.md
@@ -118,7 +118,7 @@
         } else {
           hasChild = Math.random() > 0.5;
         }
-    
+
         setTimeout(function() {
           let data;
           if (hasChild) {
@@ -130,7 +130,7 @@
           } else {
             data = [];
           }
-    
+
           resolve(data);
         }, 500);
       },
@@ -161,11 +161,11 @@
       append(store, data) {
         store.append({ id: id++, label: 'testtest', children: [] }, data);
       },
-  
+
       remove(store, data) {
         store.remove(data);
       },
-  
+
       renderContent(h, { node, data, store }) {
         return (
           <span>
@@ -178,13 +178,13 @@
             </span>
           </span>);
       },
-  
+
       filterNode(value, data) {
         if (!value) return true;
         return data.label.indexOf(value) !== -1;
       }
     },
-    
+
     data() {
       return {
         data,
@@ -200,7 +200,7 @@
   };
 </script>
 
-## Tree 
+## Tree
 
 Display a set of data with hierarchies.
 
@@ -452,7 +452,7 @@ Tree nodes can be initially expanded or checked
         this.$refs.tree.setCheckedKeys([]);
       }
     },
-    
+
     data() {
       return {
         data2: [{
@@ -518,7 +518,7 @@ The content of tree nodes can be customized, so you can add icons or buttons as
 
 <script>
   let id = 1000;
-  
+
   export default {
     data() {
       return {
@@ -563,16 +563,16 @@ The content of tree nodes can be customized, so you can add icons or buttons as
         }
       }
     },
-    
+
     methods: {
       append(store, data) {
         store.append({ id: id++, label: 'testtest', children: [] }, data);
       },
-      
+
       remove(store, data) {
         store.remove(data);
       },
-      
+
       renderContent(h, { node, data, store }) {
         return (
           <span>
@@ -777,8 +777,8 @@ Only one node among the same level can be expanded at one time.
 | filter          | filter all tree nodes, filtered nodes will be hidden | Accept a parameter which will be used as first parameter for filter-node-method |
 | getCheckedNodes | If the node can be selected (`show-checkbox` is `true`), it returns the currently selected array of nodes | Accept a boolean type parameter whose default value is `false`. If the parameter is `true`, it only returns the currently selected array of sub-nodes. |
 | setCheckedNodes | set certain nodes to be checked, only works when `node-key` is assigned | an array of nodes to be checked          |
-| getCheckedKeys  | If the node can be selected (`show-checkbox` is `true`), it returns the currently selected array of node's keys | (leafOnly) Accept a boolean type parameter whose default value is `true`. If the parameter is `true`, it only returns the currently selected array of sub-nodes. |
-| setCheckedKeys  | set certain nodes to be checked, only works when `node-key` is assigned | (keys, leafOnly) Accept two parameters: 1. an array of node's keys to be checked 2. a boolean type parameter whose default value is `true`. If the parameter is `true`, it only returns the currently selected array of sub-nodes. |
+| getCheckedKeys  | If the node can be selected (`show-checkbox` is `true`), it returns the currently selected array of node's keys | (leafOnly) Accept a boolean type parameter whose default value is `false`. If the parameter is `true`, it only returns the currently selected array of sub-nodes. |
+| setCheckedKeys  | set certain nodes to be checked, only works when `node-key` is assigned | (keys, leafOnly) Accept two parameters: 1. an array of node's keys to be checked 2. a boolean type parameter whose default value is `false`. If the parameter is `true`, it only returns the currently selected array of sub-nodes. |
 | setChecked      | set node to be checked or not, only works when `node-key` is assigned | (key/data, checked, deep) Accept three parameters: 1. node's key or data to be checked 2. a boolean typed parameter indicating checked or not. 3. a boolean typed parameter indicating deep or not. |
 
 ### Events
@@ -789,4 +789,3 @@ Only one node among the same level can be expanded at one time.
 | current-change | triggers when current node changes       | two parameters: node object corresponding to the current node, `node` property of TreeNode |
 | node-expand    | triggers when current node open          | three parameters: node object corresponding to the node opened, `node` property of TreeNode, TreeNode itself |
 | node-collapse  | triggers when current node close         | three parameters: node object corresponding to the node closed, `node` property of TreeNode, TreeNode itself |
-
diff --git a/examples/docs/zh-CN/tree.md b/examples/docs/zh-CN/tree.md
index 99346b6c4..bf1635e30 100644
--- a/examples/docs/zh-CN/tree.md
+++ b/examples/docs/zh-CN/tree.md
@@ -4,16 +4,16 @@
       width: 20px;
       background: #ddd;
     }
-  
+
     .folder {
       width: 20px;
       background: #888;
     }
-    
+
     .buttons {
       margin-top: 20px;
     }
-    
+
     .filter-tree {
       margin-top: 20px;
     }
@@ -140,7 +140,7 @@
         } else {
           hasChild = Math.random() > 0.5;
         }
-    
+
         setTimeout(function() {
           var data;
           if (hasChild) {
@@ -152,7 +152,7 @@
           } else {
             data = [];
           }
-    
+
           resolve(data);
         }, 500);
       },
@@ -187,7 +187,7 @@
       remove(store, data) {
         store.remove(data);
       },
-  
+
       renderContent(h, { node, data, store }) {
         return (
           <span>
@@ -200,13 +200,13 @@
             </span>
           </span>);
       },
-  
+
       filterNode(value, data) {
         if (!value) return true;
         return data.label.indexOf(value) !== -1;
       }
     },
-    
+
     data() {
       return {
         data,
@@ -474,7 +474,7 @@
         this.$refs.tree.setCheckedKeys([]);
       }
     },
-    
+
     data() {
       return {
         data2: [{
@@ -594,7 +594,7 @@
       remove(store, data) {
         store.remove(data);
       },
-  
+
       renderContent(h, { node, data, store }) {
         return (
           <span>
@@ -646,7 +646,7 @@
         return data.label.indexOf(value) !== -1;
       }
     },
-    
+
     data() {
       return {
         filterText: '',
@@ -799,8 +799,8 @@
 | filter          | 对树节点进行筛选操作                               | 接收一个任意类型的参数，该参数会在 filter-node-method 中作为第一个参数 |
 | getCheckedNodes | 若节点可被选择（即 `show-checkbox` 为 `true`），则返回目前被选中的节点所组成的数组 | (leafOnly) 接收一个 boolean 类型的参数，若为 `true` 则仅返回被选中的叶子节点，默认值为 `false` |
 | setCheckedNodes | 设置目前勾选的节点，使用此方法必须设置 node-key 属性          | (nodes) 接收勾选节点数据的数组                      |
-| getCheckedKeys  | 若节点可被选择（即 `show-checkbox` 为 `true`），则返回目前被选中的节点所组成的数组 | (leafOnly) 接收一个 boolean 类型的参数，若为 `true` 则仅返回被选中的叶子节点的 keys，默认值为 `true` |
-| setCheckedKeys  | 通过 keys 设置目前勾选的节点，使用此方法必须设置 node-key 属性  | (keys, leafOnly) 接收两个参数，1. 勾选节点的 key 的数组 2. boolean 类型的参数，若为 `true` 则仅设置叶子节点的选中状态，默认值为 `true` |
+| getCheckedKeys  | 若节点可被选择（即 `show-checkbox` 为 `true`），则返回目前被选中的节点所组成的数组 | (leafOnly) 接收一个 boolean 类型的参数，若为 `true` 则仅返回被选中的叶子节点的 keys，默认值为 `false` |
+| setCheckedKeys  | 通过 keys 设置目前勾选的节点，使用此方法必须设置 node-key 属性  | (keys, leafOnly) 接收两个参数，1. 勾选节点的 key 的数组 2. boolean 类型的参数，若为 `true` 则仅设置叶子节点的选中状态，默认值为 `false` |
 | setChecked      | 通过 key / data 设置某个节点的勾选状态，使用此方法必须设置 node-key 属性 | (key/data, checked, deep) 接收三个参数，1. 勾选节点的 key 或者 data 2. boolean 类型，节点是否选中  3. boolean 类型，是否设置子节点 ，默认为 false |
 
 ### Events
@@ -811,4 +811,3 @@
 | current-change | 当前选中节点变化时触发的事件 | 共两个参数，依次为：当前节点的数据，当前节点的 Node 对象          |
 | node-expand    | 节点被展开时触发的事件    | 共三个参数，依次为：传递给 `data` 属性的数组中该节点所对应的对象、节点对应的 Node、节点组件本身。 |
 | node-collapse  | 节点被关闭时触发的事件    | 共三个参数，依次为：传递给 `data` 属性的数组中该节点所对应的对象、节点对应的 Node、节点组件本身。 |
-
diff --git a/packages/tree/src/model/tree-store.js b/packages/tree/src/model/tree-store.js
index 50410ec7e..7755defbe 100644
--- a/packages/tree/src/model/tree-store.js
+++ b/packages/tree/src/model/tree-store.js
@@ -142,7 +142,7 @@ export default class TreeStore {
     delete this.nodesMap[node.key];
   }
 
-  getCheckedNodes(leafOnly) {
+  getCheckedNodes(leafOnly = false) {
     const checkedNodes = [];
     const traverse = function(node) {
       const childNodes = node.root ? node.root.childNodes : node.childNodes;
@@ -161,7 +161,7 @@ export default class TreeStore {
     return checkedNodes;
   }
 
-  getCheckedKeys(leafOnly) {
+  getCheckedKeys(leafOnly = false) {
     const key = this.key;
     const allNodes = this._getAllNodes();
     const keys = [];
@@ -187,18 +187,65 @@ export default class TreeStore {
     return allNodes;
   }
 
-  _setCheckedKeys(key, leafOnly, checkedKeys) {
+  _setCheckedKeys(key, leafOnly = false, checkedKeys) {
     const allNodes = this._getAllNodes();
+    allNodes.sort((a, b) => a.level < b.level);
 
-    allNodes.sort((a, b) => a.level > b.level ? -1 : 1);
+    const keys = Object.keys(checkedKeys);
     allNodes.forEach((node) => {
-      if (!leafOnly || (leafOnly && node.isLeaf)) {
-        node.setChecked(!!checkedKeys[(node.data || {})[key]], !this.checkStrictly);
+      let checked = keys.indexOf(node.data[key] + '') > -1;
+
+      if (!node.isLeaf) {
+        if (!this.checkStrictly) {
+          const childNodes = node.childNodes;
+
+          let all = true;
+          let none = true;
+
+          for (let i = 0, j = childNodes.length; i < j; i++) {
+            const child = childNodes[i];
+            if (child.checked !== true || child.indeterminate) {
+              all = false;
+            }
+            if (child.checked !== false || child.indeterminate) {
+              none = false;
+            }
+          }
+
+          if (all) {
+            node.setChecked(true, !this.checkStrictly);
+          } else if (!all && !none) {
+            checked = checked ? true : 'half';
+            node.setChecked(checked, !this.checkStrictly && checked === true);
+          } else if (none) {
+            node.setChecked(checked, !this.checkStrictly);
+          }
+        } else {
+          node.setChecked(checked, false);
+        }
+
+        if (leafOnly) {
+          node.setChecked(false, false);
+          const traverse = function(node) {
+            const childNodes = node.childNodes;
+
+            childNodes.forEach((child) => {
+              if (!child.isLeaf) {
+                child.setChecked(false, false);
+              }
+              traverse(child);
+            });
+          };
+
+          traverse(node);
+        }
+      } else {
+        node.setChecked(checked, false);
       }
     });
   }
 
-  setCheckedNodes(array, leafOnly = true) {
+  setCheckedNodes(array, leafOnly = false) {
     const key = this.key;
     const checkedKeys = {};
     array.forEach((item) => {
@@ -208,7 +255,7 @@ export default class TreeStore {
     this._setCheckedKeys(key, leafOnly, checkedKeys);
   }
 
-  setCheckedKeys(keys, leafOnly = true) {
+  setCheckedKeys(keys, leafOnly = false) {
     this.defaultCheckedKeys = keys;
     const key = this.key;
     const checkedKeys = {};
diff --git a/test/unit/specs/tree.spec.js b/test/unit/specs/tree.spec.js
index 8bfe74787..3ead00c3b 100644
--- a/test/unit/specs/tree.spec.js
+++ b/test/unit/specs/tree.spec.js
@@ -270,6 +270,46 @@ describe('Tree', () => {
     tree.setCheckedKeys([111]);
     expect(tree.getCheckedNodes().length).to.equal(3);
     expect(tree.getCheckedKeys().length).to.equal(3);
+
+    tree.setCheckedKeys([1]);
+    setTimeout(function() {
+      expect(tree.getCheckedNodes().length).to.equal(3);
+      expect(tree.getCheckedKeys().length).to.equal(3);
+    }, 0);
+
+    tree.setCheckedKeys([2]);
+    setTimeout(function() {
+      expect(tree.getCheckedNodes().length).to.equal(3);
+      expect(tree.getCheckedKeys().length).to.equal(3);
+    }, 0);
+
+    tree.setCheckedKeys([21]);
+    expect(tree.getCheckedNodes().length).to.equal(1);
+    expect(tree.getCheckedKeys().length).to.equal(1);
+  });
+
+  it('setCheckedKeys with checkStrictly', () => {
+    vm = getTreeVm(':props="defaultProps" checkStrictly show-checkbox node-key="id"');
+    const tree = vm.$children[0];
+    tree.setCheckedKeys([111]);
+    expect(tree.getCheckedNodes().length).to.equal(1);
+    expect(tree.getCheckedKeys().length).to.equal(1);
+
+    tree.setCheckedKeys([1]);
+    setTimeout(function() {
+      expect(tree.getCheckedNodes().length).to.equal(1);
+      expect(tree.getCheckedKeys().length).to.equal(1);
+    }, 0);
+
+    tree.setCheckedKeys([2]);
+    setTimeout(function() {
+      expect(tree.getCheckedNodes().length).to.equal(1);
+      expect(tree.getCheckedKeys().length).to.equal(1);
+    }, 0);
+
+    tree.setCheckedKeys([21, 22]);
+    expect(tree.getCheckedNodes().length).to.equal(2);
+    expect(tree.getCheckedKeys().length).to.equal(2);
   });
 
   it('method setChecked', () => {
@@ -288,8 +328,20 @@ describe('Tree', () => {
     vm = getTreeVm(':props="defaultProps" show-checkbox node-key="id"');
     const tree = vm.$children[0];
     tree.setCheckedKeys([1, 11, 111, 2], false);
-    expect(tree.getCheckedNodes().length).to.equal(6);
-    expect(tree.getCheckedKeys().length).to.equal(6);
+    setTimeout(function() {
+      expect(tree.getCheckedNodes().length).to.equal(6);
+      expect(tree.getCheckedKeys().length).to.equal(6);
+    }, 0);
+  });
+
+  it('setCheckedKeys with leafOnly=true', () => {
+    vm = getTreeVm(':props="defaultProps" show-checkbox node-key="id"');
+    const tree = vm.$children[0];
+    tree.setCheckedKeys([2], true);
+    setTimeout(function() {
+      expect(tree.getCheckedNodes().length).to.equal(2);
+      expect(tree.getCheckedKeys().length).to.equal(2);
+    }, 0);
   });
 
   it('check strictly', (done) => {