From 2dec8e70ec1ddb64cabf6689775fb3b8cdf21dc5 Mon Sep 17 00:00:00 2001
From: Zack Barett <zackbarett@hey.com>
Date: Tue, 24 May 2022 18:26:14 -0500
Subject: [PATCH] Integration filter for Area Selector (#12682)

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>
---
 .../ha-selector/ha-selector-area.ts           | 129 ++++++++++++------
 1 file changed, 88 insertions(+), 41 deletions(-)

diff --git a/src/components/ha-selector/ha-selector-area.ts b/src/components/ha-selector/ha-selector-area.ts
index f922d8d0c6..bdc8ec9b12 100644
--- a/src/components/ha-selector/ha-selector-area.ts
+++ b/src/components/ha-selector/ha-selector-area.ts
@@ -1,15 +1,24 @@
+import { UnsubscribeFunc } from "home-assistant-js-websocket";
 import { html, LitElement } from "lit";
 import { customElement, property, state } from "lit/decorators";
-import { ConfigEntry, getConfigEntries } from "../../data/config_entries";
+import memoizeOne from "memoize-one";
 import { DeviceRegistryEntry } from "../../data/device_registry";
-import { EntityRegistryEntry } from "../../data/entity_registry";
+import {
+  EntityRegistryEntry,
+  subscribeEntityRegistry,
+} from "../../data/entity_registry";
+import {
+  EntitySources,
+  fetchEntitySourcesWithCache,
+} from "../../data/entity_sources";
 import { AreaSelector } from "../../data/selector";
+import { SubscribeMixin } from "../../mixins/subscribe-mixin";
 import { HomeAssistant } from "../../types";
 import "../ha-area-picker";
 import "../ha-areas-picker";
 
 @customElement("ha-selector-area")
-export class HaAreaSelector extends LitElement {
+export class HaAreaSelector extends SubscribeMixin(LitElement) {
   @property() public hass!: HomeAssistant;
 
   @property() public selector!: AreaSelector;
@@ -20,29 +29,44 @@ export class HaAreaSelector extends LitElement {
 
   @property() public helper?: string;
 
-  @state() public _configEntries?: ConfigEntry[];
+  @state() private _entitySources?: EntitySources;
+
+  @state() private _entities?: EntityRegistryEntry[];
 
   @property({ type: Boolean }) public disabled = false;
 
   @property({ type: Boolean }) public required = true;
 
+  public hassSubscribe(): UnsubscribeFunc[] {
+    return [
+      subscribeEntityRegistry(this.hass.connection!, (entities) => {
+        this._entities = entities.filter((entity) => entity.device_id !== null);
+      }),
+    ];
+  }
+
   protected updated(changedProperties) {
-    if (changedProperties.has("selector")) {
-      const oldSelector = changedProperties.get("selector");
-      if (
-        oldSelector !== this.selector &&
-        this.selector.area.device?.integration
-      ) {
-        getConfigEntries(this.hass, {
-          domain: this.selector.area.device.integration,
-        }).then((entries) => {
-          this._configEntries = entries;
-        });
-      }
+    if (
+      changedProperties.has("selector") &&
+      (this.selector.area.device?.integration ||
+        this.selector.area.entity?.integration) &&
+      !this._entitySources
+    ) {
+      fetchEntitySourcesWithCache(this.hass).then((sources) => {
+        this._entitySources = sources;
+      });
     }
   }
 
   protected render() {
+    if (
+      (this.selector.area.device?.integration ||
+        this.selector.area.entity?.integration) &&
+      !this._entitySources
+    ) {
+      return html``;
+    }
+
     if (!this.selector.area.multiple) {
       return html`
         <ha-area-picker
@@ -87,39 +111,62 @@ export class HaAreaSelector extends LitElement {
   }
 
   private _filterEntities = (entity: EntityRegistryEntry): boolean => {
-    if (this.selector.area.entity?.integration) {
-      if (entity.platform !== this.selector.area.entity.integration) {
+    const filterIntegration = this.selector.area.entity?.integration;
+    if (
+      filterIntegration &&
+      this._entitySources?.[entity.entity_id]?.domain !== filterIntegration
+    ) {
+      return false;
+    }
+    return true;
+  };
+
+  private _filterDevices = (device: DeviceRegistryEntry): boolean => {
+    if (!this.selector.area.device) {
+      return true;
+    }
+
+    const {
+      manufacturer: filterManufacturer,
+      model: filterModel,
+      integration: filterIntegration,
+    } = this.selector.area.device;
+
+    if (filterManufacturer && device.manufacturer !== filterManufacturer) {
+      return false;
+    }
+    if (filterModel && device.model !== filterModel) {
+      return false;
+    }
+    if (filterIntegration && this._entitySources && this._entities) {
+      const deviceIntegrations = this._deviceIntegrations(
+        this._entitySources,
+        this._entities
+      );
+      if (!deviceIntegrations?.[device.id]?.includes(filterIntegration)) {
         return false;
       }
     }
     return true;
   };
 
-  private _filterDevices = (device: DeviceRegistryEntry): boolean => {
-    if (
-      this.selector.area.device?.manufacturer &&
-      device.manufacturer !== this.selector.area.device.manufacturer
-    ) {
-      return false;
-    }
-    if (
-      this.selector.area.device?.model &&
-      device.model !== this.selector.area.device.model
-    ) {
-      return false;
-    }
-    if (this.selector.area.device?.integration) {
-      if (
-        this._configEntries &&
-        !this._configEntries.some((entry) =>
-          device.config_entries.includes(entry.entry_id)
-        )
-      ) {
-        return false;
+  private _deviceIntegrations = memoizeOne(
+    (entitySources: EntitySources, entities: EntityRegistryEntry[]) => {
+      const deviceIntegrations: Record<string, string[]> = {};
+
+      for (const entity of entities) {
+        const source = entitySources[entity.entity_id];
+        if (!source?.domain) {
+          continue;
+        }
+        if (!deviceIntegrations[entity.device_id!]) {
+          deviceIntegrations[entity.device_id!] = [];
+        }
+        deviceIntegrations[entity.device_id!].push(source.domain);
       }
+      return deviceIntegrations;
     }
-    return true;
-  };
+  );
 }
 
 declare global {