Restore focus after a HTML dialog is shown modally

dialog-focus-restore.js

/**
 * Updates the passed dialog to retain focus and restore it when the dialog is closed. Won't
 * upgrade a dialog more than once. Supports IE11+ and is a no-op otherwise.
 * @param {!HTMLDialogElement} dialog to upgrade
 */
window.registerFocusRestoreDialog = (function() {
  var WeakMap = window.WeakMap;
  var registered =
    typeof WeakMap === 'function'
      ? new WeakMap()
      : (function() {
          var objKeyStr =
            '__WeakMapData.' +
            Math.random()
              .toString(36)
              .substr(2) +
            '__';
          var hasKey = function(key) {
            return Object.prototype.hasOwnProperty.call(key, objKeyStr);
          };

          return {
            get: function(key) {
              return hasKey(key) ? key[objKeyStr] : undefined;
            },
            set: function(key, value) {
              if (!hasKey(key) && typeof Object.defineProperty === 'function') {
                Object.defineProperty(key, objKeyStr, {
                  configurable: false,
                  enumerable: false,
                  value: value,
                  writable: true
                });
              }
              key[objKeyStr] = value;
            },
            has: hasKey
          };
        })();

  var focusoutListener = function(ev) {
    var document = ev.currentTarget;
    var previousFocus = ev.target;
    registered.set(document, previousFocus);
  };

  var closeListener = function(ev) {
    var dialog = ev.currentTarget;

    if (dialog.hasAttribute('open')) {
      return; // in native, this fires the frame later
    }

    focusOnSaved(dialog);
  };

  var focusOnSaved = function(dialog) {
    var document = dialog.ownerDocument || document;

    var savedFocus = registered.get(dialog);
    if (document.contains(savedFocus)) {
      var wasFocus = document.activeElement;
      savedFocus.focus();
      if (document.activeElement !== savedFocus) {
        wasFocus.focus(); // restore focus, we couldn't focus saved
      }
    }
    savedFocus = null;
  };

  var mo = null;
  if (window.MutationObserver) {
    // watch for 'open' change and clear saved
    mo = new MutationObserver(function(records) {
      for (var i = records.length; i--; ) {
        var record = records[i];
        var dialog = record.target;

        if (!dialog.hasAttribute('open')) {
          focusOnSaved(dialog);
          registered.set(dialog, null);
        } else {
          // if open was cleared/set in the same frame, then the dialog will still be a modal (Y)
        }
      }
    });
  }

  return function registerFocusRestoreDialog(dialog) {
    if (dialog.tagName !== 'DIALOG') {
      throw new Error(
        'Failed to upgrade focus on dialog: The element is not a dialog.'
      );
    }

    var document = dialog.ownerDocument || document;
    if (!registered.has(document)) {
      // store previous focused node centrally
      registered.set(document, null);
      document.addEventListener('focusout', focusoutListener, true);
    }

    if (registered.has(dialog)) {
      return;
    }
    registered.set(dialog, null);

    // replace showModal method directly, to save focus
    var realShowModal = dialog.showModal;
    dialog.showModal = function showModal(anchor) {
      var savedFocus = document.activeElement;
      if (savedFocus === document || savedFocus === document.body) {
        // some browsers read activeElement as body
        var previousFocus = registered.get(document);
        savedFocus = previousFocus;
      }

      registered.set(dialog, savedFocus);

      if (arguments.length < 1) {
        realShowModal.call(this);
      } else {
        realShowModal.call(this, anchor);
      }
    };

    if (mo) {
      // watch for 'open' change and clear saved
      mo.observe(dialog, { attributes: true, attributeFilter: ['open'] });
    }

    // on close, try to focus saved, if possible
    dialog.addEventListener('close', closeListener);

    // FIXME: If a modal dialog is readded to the page (either remove/add or .appendChild), it will
    // be a non-modal. It will still have its 'close' handler called and try to focus on the saved
    // element.
    //
    // These could basically be solved if 'close' yielded whether it was a modal or non-modal
    // being closed. But it doesn't. It could also be solved by a permanent MutationObserver, as is
    // done inside the polyfill.
  };
})();