From bff442d3e188a7ced071230ff45ab4452cd85b42 Mon Sep 17 00:00:00 2001
From: Mingze Xiao <mxiao@box.com>
Date: Tue, 16 Jun 2020 15:09:12 -0700
Subject: [PATCH] feat(annotations): Address comments

---
 src/lib/AnnotationControls.ts        | 27 +++++++++++++++++++++++++--
 src/lib/viewers/BaseViewer.js        | 16 ++--------------
 src/lib/viewers/doc/DocBaseViewer.js |  4 +---
 3 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/src/lib/AnnotationControls.ts b/src/lib/AnnotationControls.ts
index efdbed976..aada93bb1 100644
--- a/src/lib/AnnotationControls.ts
+++ b/src/lib/AnnotationControls.ts
@@ -25,6 +25,10 @@ interface ControlsMap {
     [key: string]: () => void;
 }
 
+const buttonClassMap = {
+    region: CLASS_REGION_BUTTON,
+};
+
 export default class AnnotationControls {
     private controls: Controls;
 
@@ -60,7 +64,11 @@ export default class AnnotationControls {
     }
 
     public focus(): void {
-        this.controlsElement.querySelector(`.${CLASS_REGION_BUTTON}`).focus();
+        const buttonClass = buttonClassMap[this.currentActiveControl];
+        if (!buttonClass) {
+            return;
+        }
+        this.controlsElement.querySelector(`.${buttonClass}`).focus();
     }
 
     /**
@@ -152,7 +160,7 @@ export default class AnnotationControls {
     /**
      * Initialize the annotation controls with options.
      */
-    public init({ onRegionClick = noop }: Options = {}): void {
+    public init({ onRegionClick = noop, annotator }: Options = {}): void {
         const groupElement = this.controls.addGroup(CLASS_ANNOTATIONS_GROUP);
         const regionButton = this.controls.add(
             __('region_comment'),
@@ -164,5 +172,20 @@ export default class AnnotationControls {
         );
 
         regionButton.setAttribute('data-testid', 'bp-AnnotationsControls-regionBtn');
+        regionButton.onkeydown = (event: KeyboardEvent): void => {
+            if (!annotator) {
+                return;
+            }
+
+            if (event.key !== 'Escape') {
+                return;
+            }
+
+            this.resetControls();
+            annotator.toggleAnnotationMode(AnnotationMode.NONE);
+
+            event.preventDefault();
+            event.stopPropagation();
+        };
     }
 }
diff --git a/src/lib/viewers/BaseViewer.js b/src/lib/viewers/BaseViewer.js
index e4d8fca03..973a5a7ed 100644
--- a/src/lib/viewers/BaseViewer.js
+++ b/src/lib/viewers/BaseViewer.js
@@ -1178,8 +1178,8 @@ class BaseViewer extends EventEmitter {
         }
     }
 
-    handleAnnotationStagedEvent(staged) {
-        if (!staged && this.annotationControls) {
+    handleAnnotationStagedEvent({ staged }) {
+        if (this.annotationControls && !staged) {
             this.annotationControls.focus();
         }
     }
@@ -1264,18 +1264,6 @@ class BaseViewer extends EventEmitter {
 
         return addedElement;
     }
-
-    /**
-     * Method to exit annotation mode
-     *
-     * @return {void}
-     */
-    exitAnnotationMode = () => {
-        if (this.areAnnotationsEnabled() && this.annotator && this.annotationControls) {
-            this.annotator.toggleAnnotationMode(AnnotationMode.NONE);
-            this.annotationControls.resetControls();
-        }
-    };
 }
 
 export default BaseViewer;
diff --git a/src/lib/viewers/doc/DocBaseViewer.js b/src/lib/viewers/doc/DocBaseViewer.js
index 75711aad9..176334148 100644
--- a/src/lib/viewers/doc/DocBaseViewer.js
+++ b/src/lib/viewers/doc/DocBaseViewer.js
@@ -593,9 +593,6 @@ class DocBaseViewer extends BaseViewer {
             case ']':
                 this.nextPage();
                 break;
-            case 'Escape':
-                this.exitAnnotationMode();
-                break;
             default:
                 if (this.findBar) {
                     return this.findBar.onKeydown(event);
@@ -1105,6 +1102,7 @@ class DocBaseViewer extends BaseViewer {
 
         if (this.areNewAnnotationsEnabled() && this.hasAnnotationCreatePermission()) {
             this.annotationControls.init({
+                annotator: this.annotator,
                 onRegionClick: this.regionClickHandler,
             });
         }