MDL-72523 javascript: Introduce setReturnElement for core/modal

lib/amd/src/modal.js

@@ -103,6 +103,7 @@ define([
         this.modalCount = modalCounter++;
         this.attachmentPoint = document.createElement('div');
         document.body.append(this.attachmentPoint);
+        this.focusOnClose = null;
 
         if (!this.root.is(SELECTORS.CONTAINER)) {
             Notification.exception({message: 'Element is not a modal container'});
@@ -693,6 +694,11 @@ define([
 
         this.attachToDOM();
 
+        // If the focusOnClose was not set. Set the focus back to triggered element.
+        if (!this.focusOnClose && document.activeElement) {
+            this.focusOnClose = document.activeElement;
+        }
+
         return this.getBackdrop()
         .then(function(backdrop) {
             var currentIndex = this.calculateZIndex();
@@ -850,6 +856,13 @@ define([
             this.hide();
             data.originalEvent.preventDefault();
         }.bind(this));
+
+        this.getRoot().on(ModalEvents.hidden, () => {
+            if (this.focusOnClose) {
+                // Focus on the element that actually triggers the modal.
+                this.focusOnClose.focus();
+            }
+        });
     };
 
     /**
@@ -962,5 +975,14 @@ define([
         this.removeOnClose = remove;
     };
 
+    /**
+     * Set the return element for the modal.
+     *
+     * @param {Element|jQuery} element Element to focus when the modal is closed
+     */
+    Modal.prototype.setReturnElement = function(element) {
+        this.focusOnClose = element;
+    };
+
     return Modal;
 });