fix(dialog): capture previously focused element immediately

Captures the previously-focused element immediately, instead of waiting until the animation is done. This fixes a regression from #3774, because if we wait for the animation to finish, the focus might have shifted, or the element could have been disabled.

Author: crisbeto

src/lib/dialog/dialog-container.ts

@@ -8,6 +8,8 @@ import {
   Renderer,
   ElementRef,
   EventEmitter,
+  Inject,
+  Optional,
 } from '@angular/core';
 import {
   animate,
@@ -17,6 +19,7 @@ import {
   transition,
   AnimationEvent,
 } from '@angular/animations';
+import {DOCUMENT} from '@angular/platform-browser';
 import {BasePortalHost, ComponentPortal, PortalHostDirective, TemplatePortal} from '../core';
 import {MdDialogConfig} from './dialog-config';
 import {MdDialogContentAlreadyAttachedError} from './dialog-errors';
@@ -64,6 +67,9 @@ export class MdDialogContainer extends BasePortalHost implements OnDestroy {
   /** Element that was focused before the dialog was opened. Save this to restore upon close. */
   private _elementFocusedBeforeDialogWasOpened: HTMLElement = null;
 
+  /** Reference to the global document object. */
+  private _document: Document|Document;
+
   /** The dialog configuration. */
   dialogConfig: MdDialogConfig;
 
@@ -77,9 +83,11 @@ export class MdDialogContainer extends BasePortalHost implements OnDestroy {
     private _ngZone: NgZone,
     private _renderer: Renderer,
     private _elementRef: ElementRef,
-    private _focusTrapFactory: FocusTrapFactory) {
+    private _focusTrapFactory: FocusTrapFactory,
+    @Optional() @Inject(DOCUMENT) _document: any) {
 
     super();
+    this._document = _document;
   }
 
   /**
@@ -91,6 +99,7 @@ export class MdDialogContainer extends BasePortalHost implements OnDestroy {
       throw new MdDialogContentAlreadyAttachedError();
     }
 
+    this._savePreviouslyFocusedElement();
     return this._portalHost.attachComponentPortal(portal);
   }
 
@@ -103,12 +112,12 @@ export class MdDialogContainer extends BasePortalHost implements OnDestroy {
       throw new MdDialogContentAlreadyAttachedError();
     }
 
+    this._savePreviouslyFocusedElement();
     return this._portalHost.attachTemplatePortal(portal);
   }
 
   /**
    * Moves the focus inside the focus trap.
-   * @private
    */
   private _trapFocus() {
     if (!this._focusTrap) {
@@ -118,10 +127,18 @@ export class MdDialogContainer extends BasePortalHost implements OnDestroy {
     // If were to attempt to focus immediately, then the content of the dialog would not yet be
     // ready in instances where change detection has to run first. To deal with this, we simply
     // wait for the microtask queue to be empty.
-    this._elementFocusedBeforeDialogWasOpened = document.activeElement as HTMLElement;
     this._focusTrap.focusFirstTabbableElementWhenReady();
   }
 
+  /**
+   * Saves a reference to the element that was focused before the dialog was opened.
+   */
+  private _savePreviouslyFocusedElement() {
+    if (this._document) {
+      this._elementFocusedBeforeDialogWasOpened = this._document.activeElement as HTMLElement;
+    }
+  }
+
   /**
    * Kicks off the leave animation.
    * @docs-private