Return focus after dialogs close (#11999)

src/common/dom/ancestors-with-property.ts

@@ -0,0 +1,41 @@
+const DEFAULT_OWN = true;
+
+// Finds the closest ancestor of an element that has a specific optionally owned property,
+// traversing slot and shadow root boundaries until the body element is reached
+export const closestWithProperty = (
+  element: Element | null,
+  property: string | symbol,
+  own = DEFAULT_OWN
+) => {
+  if (!element || element === document.body) return null;
+
+  element = element.assignedSlot ?? element;
+  if (element.parentElement) {
+    element = element.parentElement;
+  } else {
+    const root = element.getRootNode();
+    element = root instanceof ShadowRoot ? root.host : null;
+  }
+
+  if (
+    own
+      ? Object.prototype.hasOwnProperty.call(element, property)
+      : element && property in element
+  )
+    return element;
+  return closestWithProperty(element, property, own);
+};
+
+// Finds the set of all such ancestors and includes starting element as first in the set
+export const ancestorsWithProperty = (
+  element: Element | null,
+  property: string | symbol,
+  own = DEFAULT_OWN
+) => {
+  const ancestors: Set<Element> = new Set();
+  while (element) {
+    ancestors.add(element);
+    element = closestWithProperty(element, property, own);
+  }
+  return ancestors;
+};

src/components/ha-button-menu.ts

@@ -1,17 +1,27 @@
+import type { Button } from "@material/mwc-button";
 import "@material/mwc-menu";
 import type { Corner, Menu, MenuCorner } from "@material/mwc-menu";
 import { css, CSSResultGroup, html, LitElement, TemplateResult } from "lit";
-import { customElement, property, query } from "lit/decorators";
+import {
+  customElement,
+  property,
+  query,
+  queryAssignedElements,
+} from "lit/decorators";
+import { FOCUS_TARGET } from "../dialogs/make-dialog-manager";
+import type { HaIconButton } from "./ha-icon-button";
 
 @customElement("ha-button-menu")
 export class HaButtonMenu extends LitElement {
+  protected readonly [FOCUS_TARGET];
+
   @property() public corner: Corner = "TOP_START";
 
   @property() public menuCorner: MenuCorner = "START";
 
-  @property({ type: Number }) public x?: number;
+  @property({ type: Number }) public x: number | null = null;
 
-  @property({ type: Number }) public y?: number;
+  @property({ type: Number }) public y: number | null = null;
 
   @property({ type: Boolean }) public multi = false;
 
@@ -23,6 +33,12 @@ export class HaButtonMenu extends LitElement {
 
   @query("mwc-menu", true) private _menu?: Menu;
 
+  @queryAssignedElements({
+    slot: "trigger",
+    selector: "ha-icon-button, mwc-button",
+  })
+  private _triggerButton!: Array<HaIconButton | Button>;
+
   public get items() {
     return this._menu?.items;
   }
@@ -31,6 +47,14 @@ export class HaButtonMenu extends LitElement {
     return this._menu?.selected;
   }
 
+  public override focus() {
+    if (this._menu?.open) {
+      this._menu.focusItemAtIndex(0);
+    } else {
+      this._triggerButton[0]?.focus();
+    }
+  }
+
   protected render(): TemplateResult {
     return html`
       <div @click=${this._handleClick}>

src/components/ha-dialog.ts

@@ -4,6 +4,7 @@ import { mdiClose } from "@mdi/js";
 import { css, html, TemplateResult } from "lit";
 import { customElement } from "lit/decorators";
 import type { HomeAssistant } from "../types";
+import { FOCUS_TARGET } from "../dialogs/make-dialog-manager";
 import "./ha-icon-button";
 
 export const createCloseHeading = (
@@ -21,6 +22,8 @@ export const createCloseHeading = (
 
 @customElement("ha-dialog")
 export class HaDialog extends DialogBase {
+  protected readonly [FOCUS_TARGET];
+
   public scrollToPos(x: number, y: number) {
     this.contentElement?.scrollTo(x, y);
   }

src/components/ha-icon-button.ts

@@ -1,6 +1,7 @@
 import "@material/mwc-icon-button";
+import type { IconButton } from "@material/mwc-icon-button";
 import { css, CSSResultGroup, html, LitElement, TemplateResult } from "lit";
-import { customElement, property } from "lit/decorators";
+import { customElement, property, query } from "lit/decorators";
 import "./ha-svg-icon";
 
 @customElement("ha-icon-button")
@@ -15,6 +16,12 @@ export class HaIconButton extends LitElement {
 
   @property({ type: Boolean }) hideTitle = false;
 
+  @query("mwc-icon-button", true) private _button?: IconButton;
+
+  public override focus() {
+    this._button?.focus();
+  }
+
   static shadowRootOptions: ShadowRootInit = {
     mode: "open",
     delegatesFocus: true,

src/dialogs/make-dialog-manager.ts

@@ -1,6 +1,9 @@
 import { HASSDomEvent, ValidHassDomEvent } from "../common/dom/fire_event";
 import { mainWindow } from "../common/dom/get_main_window";
 import { ProvideHassElement } from "../mixins/provide-hass-lit-mixin";
+import { ancestorsWithProperty } from "../common/dom/ancestors-with-property";
+import { deepActiveElement } from "../common/dom/deep-active-element";
+import { nextRender } from "../common/util/render-status";
 
 declare global {
   // for fire event
@@ -40,7 +43,17 @@ export interface DialogState {
   dialogParams?: unknown;
 }
 
-const LOADED = {};
+interface LoadedDialogInfo {
+  element: Promise<HassDialog>;
+  closedFocusTargets?: Set<Element>;
+}
+
+interface LoadedDialogsDict {
+  [tag: string]: LoadedDialogInfo;
+}
+
+const LOADED: LoadedDialogsDict = {};
+export const FOCUS_TARGET = Symbol.for("HA focus target");
 
 export const showDialog = async (
   element: HTMLElement & ProvideHassElement,
@@ -60,11 +73,24 @@ export const showDialog = async (
       }
       return;
     }
-    LOADED[dialogTag] = dialogImport().then(() => {
+    LOADED[dialogTag] = {
+      element: dialogImport().then(() => {
         const dialogEl = document.createElement(dialogTag) as HassDialog;
         element.provideHass(dialogEl);
         return dialogEl;
-    });
+      }),
+    };
+  }
+
+  // Get the focus targets after the dialog closes, but keep the original if dialog is being replaced
+  if (mainWindow.history.state?.replaced) {
+    LOADED[dialogTag].closedFocusTargets =
+      LOADED[mainWindow.history.state.dialog].closedFocusTargets;
+  } else {
+    LOADED[dialogTag].closedFocusTargets = ancestorsWithProperty(
+      deepActiveElement(),
+      FOCUS_TARGET
+    );
   }
 
   if (addHistory) {
@@ -93,25 +119,29 @@ export const showDialog = async (
       );
     }
   }
-  const dialogElement = await LOADED[dialogTag];
+
+  const dialogElement = await LOADED[dialogTag].element;
+  dialogElement.addEventListener("dialog-closed", _handleClosedFocus);
+
   // Append it again so it's the last element in the root,
   // so it's guaranteed to be on top of the other elements
   root.appendChild(dialogElement);
   dialogElement.showDialog(dialogParams);
 };
 
-export const replaceDialog = () => {
+export const replaceDialog = (dialogElement: HassDialog) => {
   mainWindow.history.replaceState(
     { ...mainWindow.history.state, replaced: true },
     ""
   );
+  dialogElement.removeEventListener("dialog-closed", _handleClosedFocus);
 };
 
 export const closeDialog = async (dialogTag: string): Promise<boolean> => {
   if (!(dialogTag in LOADED)) {
     return true;
   }
-  const dialogElement: HassDialog = await LOADED[dialogTag];
+  const dialogElement = await LOADED[dialogTag].element;
   if (dialogElement.closeDialog) {
     return dialogElement.closeDialog() !== false;
   }
@@ -137,3 +167,33 @@ export const makeDialogManager = (
     }
   );
 };
+
+const _handleClosedFocus = async (ev: HASSDomEvent<DialogClosedParams>) => {
+  const closedFocusTargets = LOADED[ev.detail.dialog].closedFocusTargets;
+  delete LOADED[ev.detail.dialog].closedFocusTargets;
+  if (!closedFocusTargets) return;
+
+  // Undo whatever the browser focused to provide easy checking
+  let focusedElement = deepActiveElement();
+  if (focusedElement instanceof HTMLElement) focusedElement.blur();
+
+  // Make sure backdrop is fully updated before trying (especially needed for underlay dialogs)
+  await nextRender();
+
+  // Try all targets in order and stop when one works
+  for (const focusTarget of closedFocusTargets) {
+    if (focusTarget instanceof HTMLElement) {
+      focusTarget.focus();
+      focusedElement = deepActiveElement();
+      if (focusedElement && focusedElement !== document.body) return;
+    }
+  }
+
+  if (__DEV__) {
+    // eslint-disable-next-line no-console
+    console.warn(
+      "Failed to focus any targets after closing dialog: %o",
+      closedFocusTargets
+    );
+  }
+};

src/dialogs/more-info/ha-more-info-dialog.ts

@@ -295,7 +295,7 @@ export class MoreInfoDialog extends LitElement {
   }
 
   private _gotoSettings() {
-    replaceDialog();
+    replaceDialog(this);
     showEntityEditorDialog(this, {
       entity_id: this._entityId!,
     });

src/panels/config/entities/dialog-entity-editor.ts

@@ -220,7 +220,7 @@ export class DialogEntityEditor extends LitElement {
   }
 
   private _openMoreInfo(): void {
-    replaceDialog();
+    replaceDialog(this);
     fireEvent(this, "hass-more-info", {
       entityId: this._params!.entity_id,
     });