From 2466ea61b08d42ba909ea9e70d98e624475f0b66 Mon Sep 17 00:00:00 2001
From: Adam L <seethroughtrees@gmail.com>
Date: Mon, 26 Mar 2018 00:25:24 -0700
Subject: [PATCH] refactor findOne method

---
 konva.d.ts                  |  2 +-
 src/Container.js            | 43 ++++++++++++++++++++++++++-----------
 test/unit/Container-test.js |  6 +++++-
 3 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/konva.d.ts b/konva.d.ts
index 24982724..85361766 100644
--- a/konva.d.ts
+++ b/konva.d.ts
@@ -343,7 +343,7 @@ declare namespace Konva {
     clipFunc(ctx: CanvasRenderingContext2D | undefined | null): void;
     destroyChildren(): void;
     find(selector?: string | ((Node) => boolean)): Collection;
-    findOne<T extends Node>(selector: string): T;
+    findOne<T extends Node>(selector: string | ((Node) => boolean)): T;
     getAllIntersections(pos: Vector2d): Shape[];
     hasChildren(): boolean;
     removeChildren(): void;
diff --git a/src/Container.js b/src/Container.js
index d7abb147..86db3ae4 100644
--- a/src/Container.js
+++ b/src/Container.js
@@ -181,31 +181,42 @@
      * });
      */
     find: function(selector) {
-      var retArr = [];
-
-      if (typeof selector === 'string') {
-        retArr = this._findByString(selector);
-      } else if (typeof selector === 'function') {
-        retArr = this._findByFunction(selector);
-      }
-
-      return Konva.Collection.toCollection(retArr);
+      // protecting _findByType to prevent user from accidentally adding
+      // second argument and getting unexpected `findOne` result
+      return this._findByType(selector, false);
     },
     /**
      * return a first node from `find` method
      * @method
      * @memberof Konva.Container.prototype
-     * @param {String} selector
-     * @returns {Konva.Node}
+     * @param {String | Function} selector
+     * @returns {Konva.Node | Undefined}
      * @example
      * // select node with id foo
      * var node = stage.findOne('#foo');
      *
      * // select node with name bar inside layer
      * var nodes = layer.findOne('.bar');
+     *
+     * // select the first node to return true in a function
+     * var node = stage.findOne(node => {
+     *  return node.getType() === 'Shape'
+     * })
      */
     findOne: function(selector) {
-      return this.find(selector)[0];
+      var result = this._findByType(selector, true);
+      return result.length > 0 ? result[0] : undefined;
+    },
+    _findByType: function(selector, findOne) {
+      var retArr = [];
+
+      if (typeof selector === 'string') {
+        retArr = this._findByString(selector, findOne);
+      } else if (typeof selector === 'function') {
+        retArr = this._findByFunction(selector, findOne);
+      }
+
+      return Konva.Collection.toCollection(retArr);
     },
     _findByString: function(selector) {
       var retArr = [],
@@ -254,10 +265,16 @@
 
       return retArr;
     },
-    _findByFunction: function(fn) {
+    // (fn: ((Node) => boolean, findOne?: boolean)
+    _findByFunction: function(fn, findOne) {
       var retArr = [];
 
       var addItems = function(el) {
+        // escape function if we've already found one.
+        if (findOne && retArr.length > 0) {
+          return;
+        }
+
         var children = el.getChildren();
         var clen = children.length;
 
diff --git a/test/unit/Container-test.js b/test/unit/Container-test.js
index 14c27169..994aabb6 100644
--- a/test/unit/Container-test.js
+++ b/test/unit/Container-test.js
@@ -291,7 +291,7 @@ suite('Container', function() {
   });
 
   // ======================================================
-  test('select shape by id and name with findOne', function() {
+  test('select shape with findOne', function() {
     var stage = addStage();
     var layer = new Konva.Layer({
       id: 'myLayer'
@@ -330,6 +330,10 @@ suite('Container', function() {
     assert.equal(node, undefined, 'node should be undefined');
     node = stage.findOne('#myLayer');
     assert.equal(node, layer, 'node type should be Layer');
+    node = stage.findOne(function(node) {
+      return node.getType() === 'Shape';
+    });
+    assert.equal(node, circle, 'findOne should work with functions');
   });
 
   // ======================================================