From 31402a3b8316f5fbbe1588c7142227e3081599e1 Mon Sep 17 00:00:00 2001
From: Max Okorokov <max@okorokov.com>
Date: Fri, 10 Jan 2020 17:41:22 +0100
Subject: [PATCH] feat(datepicker): add 'restoreFocus' input (#3539)

restoreFocus input will define re-focusing strategy when closing the input datepicker.

```
@input() restoreFocus: true | string | HTMLElement; // defaults to true
```

By default it re-focuses element focused at the moment of datepicker opening.
Otherwise accepts either selector or element reference.

Fixes #3483
---
 .../datepicker-input-config.spec.ts           |  1 +
 src/datepicker/datepicker-input-config.ts     |  1 +
 src/datepicker/datepicker-input.spec.ts       | 98 +++++++++++++++++++
 src/datepicker/datepicker-input.ts            | 27 ++++-
 4 files changed, 123 insertions(+), 4 deletions(-)

diff --git a/src/datepicker/datepicker-input-config.spec.ts b/src/datepicker/datepicker-input-config.spec.ts
index 7df4346ada..b37b06ed4a 100644
--- a/src/datepicker/datepicker-input-config.spec.ts
+++ b/src/datepicker/datepicker-input-config.spec.ts
@@ -8,5 +8,6 @@ describe('NgbInputDatepickerConfig', () => {
     expect(config.container).toBeUndefined();
     expect(config.positionTarget).toBeUndefined();
     expect(config.placement).toEqual(['bottom-left', 'bottom-right', 'top-left', 'top-right']);
+    expect(config.restoreFocus).toBe(true);
   });
 });
diff --git a/src/datepicker/datepicker-input-config.ts b/src/datepicker/datepicker-input-config.ts
index 3ed85c6b31..7392ff6d65 100644
--- a/src/datepicker/datepicker-input-config.ts
+++ b/src/datepicker/datepicker-input-config.ts
@@ -15,4 +15,5 @@ export class NgbInputDatepickerConfig extends NgbDatepickerConfig {
   container: null | 'body';
   positionTarget: string | HTMLElement;
   placement: PlacementArray = ['bottom-left', 'bottom-right', 'top-left', 'top-right'];
+  restoreFocus: true | HTMLElement | string = true;
 }
diff --git a/src/datepicker/datepicker-input.spec.ts b/src/datepicker/datepicker-input.spec.ts
index 437b87986c..90078d1402 100644
--- a/src/datepicker/datepicker-input.spec.ts
+++ b/src/datepicker/datepicker-input.spec.ts
@@ -910,6 +910,104 @@ describe('NgbInputDatepicker', () => {
        }));
   });
 
+  describe('focus restore', () => {
+
+    function open(fixture: ComponentFixture<TestComponent>) {
+      const dp = fixture.debugElement.query(By.directive(NgbInputDatepicker)).injector.get(NgbInputDatepicker);
+      dp.open();
+      fixture.detectChanges();
+    }
+
+    function selectDateAndClose(fixture: ComponentFixture<TestComponent>) {
+      fixture.nativeElement.querySelectorAll('.ngb-dp-day')[3].click();  // 1 MAR 2018
+      fixture.detectChanges();
+    }
+
+    it('should focus previously focused element', () => {
+      const fixture = createTestCmpt(`
+          <div tabindex="0" id="focusable"></div>
+          <input ngbDatepicker [startDate]="{year: 2018, month: 3}"/>
+        `);
+
+      // initial focus
+      const focusableEl = fixture.nativeElement.querySelector('#focusable');
+      focusableEl.focus();
+      expect(document.activeElement).toBe(focusableEl);
+
+      open(fixture);
+      expect(document.activeElement).not.toBe(focusableEl);
+
+      selectDateAndClose(fixture);
+      expect(document.activeElement).toBe(focusableEl);
+    });
+
+    it('should focus using selector provided via [restoreFocus]', () => {
+      const fixture = createTestCmpt(`
+          <div tabindex="0" id="focusable"></div>
+          <input ngbDatepicker restoreFocus="#focusable" [startDate]="{year: 2018, month: 3}"/>
+        `);
+
+      const focusableEl = fixture.nativeElement.querySelector('#focusable');
+      expect(document.activeElement).not.toBe(focusableEl);
+
+      open(fixture);
+      expect(document.activeElement).not.toBe(focusableEl);
+
+      selectDateAndClose(fixture);
+      expect(document.activeElement).toBe(focusableEl);
+    });
+
+    it('should focus using element provided via [restoreFocus]', () => {
+      const fixture = createTestCmpt(`
+          <div #el tabindex="0" id="focusable"></div>
+          <input ngbDatepicker [restoreFocus]="el" [startDate]="{year: 2018, month: 3}"/>
+        `);
+
+      const focusableEl = fixture.nativeElement.querySelector('#focusable');
+      expect(document.activeElement).not.toBe(focusableEl);
+
+      open(fixture);
+      expect(document.activeElement).not.toBe(focusableEl);
+
+      selectDateAndClose(fixture);
+      expect(document.activeElement).toBe(focusableEl);
+    });
+
+    it('should fallback to body if [restoreFocus] selector is invalid', () => {
+      const fixture = createTestCmpt(`
+          <div tabindex="0" id="focusable"></div>
+          <input ngbDatepicker restoreFocus=".invalid-element" [startDate]="{year: 2018, month: 3}"/>
+        `);
+
+      const focusableEl = fixture.nativeElement.querySelector('#focusable');
+      focusableEl.focus();
+      expect(document.activeElement).toBe(focusableEl);
+
+      open(fixture);
+      expect(document.activeElement).not.toBe(focusableEl);
+
+      selectDateAndClose(fixture);
+      expect(document.activeElement).toBe(document.body);
+    });
+
+    it('should fallback to body if [restoreFocus] value is invalid', () => {
+      const fixture = createTestCmpt(`
+          <div tabindex="0" id="focusable"></div>
+          <input ngbDatepicker [restoreFocus]="null" [startDate]="{year: 2018, month: 3}"/>
+        `);
+
+      const focusableEl = fixture.nativeElement.querySelector('#focusable');
+      focusableEl.focus();
+      expect(document.activeElement).toBe(focusableEl);
+
+      open(fixture);
+      expect(document.activeElement).not.toBe(focusableEl);
+
+      selectDateAndClose(fixture);
+      expect(document.activeElement).toBe(document.body);
+    });
+  });
+
   describe('Native adapter', () => {
 
     beforeEach(() => {
diff --git a/src/datepicker/datepicker-input.ts b/src/datepicker/datepicker-input.ts
index 971b7f0abd..fa0615ab3c 100644
--- a/src/datepicker/datepicker-input.ts
+++ b/src/datepicker/datepicker-input.ts
@@ -33,6 +33,7 @@ import {NgbDateParserFormatter} from './ngb-date-parser-formatter';
 import {NgbDateStruct} from './ngb-date-struct';
 import {NgbInputDatepickerConfig} from './datepicker-input-config';
 import {NgbDatepickerConfig} from './datepicker-config';
+import {isString} from '../util/util';
 
 const NGB_DATEPICKER_VALUE_ACCESSOR = {
   provide: NG_VALUE_ACCESSOR,
@@ -183,6 +184,14 @@ export class NgbInputDatepicker implements OnChanges,
    */
   @Input() placement: PlacementArray;
 
+  /**
+   * If `true`, when closing datepicker will focus element that was focused before datepicker was opened.
+   *
+   * Alternatively you could provide a selector or an `HTMLElement` to focus. If the element doesn't exist or invalid,
+   * we'll fallback to focus document body.
+   */
+  @Input() restoreFocus: true | string | HTMLElement;
+
   /**
    * If `true`, weekdays will be displayed.
    */
@@ -373,8 +382,18 @@ export class NgbInputDatepicker implements OnChanges,
       this._changeDetector.markForCheck();
 
       // restore focus
-      const elementToFocus = this._elWithFocus && this._elWithFocus['focus'] ? this._elWithFocus : this._document.body;
-      elementToFocus.focus();
+      let elementToFocus = this._elWithFocus;
+      if (isString(this.restoreFocus)) {
+        elementToFocus = this._document.querySelector(this.restoreFocus);
+      } else if (this.restoreFocus !== undefined) {
+        elementToFocus = this.restoreFocus;
+      }
+
+      if (elementToFocus) {
+        elementToFocus.focus();
+      } else {
+        this._document.body.focus();
+      }
     }
   }
 
@@ -469,8 +488,8 @@ export class NgbInputDatepicker implements OnChanges,
     }
 
     let hostElement: HTMLElement;
-    if (typeof this.positionTarget === 'string') {
-      hostElement = window.document.querySelector(this.positionTarget);
+    if (isString(this.positionTarget)) {
+      hostElement = this._document.querySelector(this.positionTarget);
     } else if (this.positionTarget instanceof HTMLElement) {
       hostElement = this.positionTarget;
     } else {