Add entity source API (#12149)

src/common/util/time-cache-entity-promise-func.ts

@@ -0,0 +1,53 @@
+import { HomeAssistant } from "../../types";
+
+interface ResultCache<T> {
+  [entityId: string]: Promise<T> | undefined;
+}
+
+/**
+ * Call a function with result caching per entity.
+ * @param cacheKey key to store the cache on hass object
+ * @param cacheTime time to cache the results
+ * @param func function to fetch the data
+ * @param hass Home Assistant object
+ * @param entityId entity to fetch data for
+ * @param args extra arguments to pass to the function to fetch the data
+ * @returns
+ */
+export const timeCacheEntityPromiseFunc = async <T>(
+  cacheKey: string,
+  cacheTime: number,
+  func: (hass: HomeAssistant, entityId: string, ...args: any[]) => Promise<T>,
+  hass: HomeAssistant,
+  entityId: string,
+  ...args: any[]
+): Promise<T> => {
+  let cache: ResultCache<T> | undefined = (hass as any)[cacheKey];
+
+  if (!cache) {
+    cache = hass[cacheKey] = {};
+  }
+
+  const lastResult = cache[entityId];
+
+  if (lastResult) {
+    return lastResult;
+  }
+
+  const result = func(hass, entityId, ...args);
+  cache[entityId] = result;
+
+  result.then(
+    // When successful, set timer to clear cache
+    () =>
+      setTimeout(() => {
+        cache![entityId] = undefined;
+      }, cacheTime),
+    // On failure, clear cache right away
+    () => {
+      cache![entityId] = undefined;
+    }
+  );
+
+  return result;
+};

src/common/util/time-cache-function-promise.ts

@@ -1,43 +1,80 @@
 import { HomeAssistant } from "../../types";
 
-interface ResultCache<T> {
-  [entityId: string]: Promise<T> | undefined;
+interface CacheResult<T> {
+  result: T;
+  cacheKey: any;
 }
 
+/**
+ * Caches a result of a promise for X time. Allows optional extra validation
+ * check to invalidate the cache.
+ * @param cacheKey the key to store the cache
+ * @param cacheTime the time to cache the result
+ * @param func the function to fetch the data
+ * @param generateCacheKey optional function to generate a cache key based on current hass + cached result. Cache is invalid if generates a different cache key.
+ * @param hass Home Assistant object
+ * @param args extra arguments to pass to the function to fetch the data
+ * @returns
+ */
 export const timeCachePromiseFunc = async <T>(
   cacheKey: string,
   cacheTime: number,
-  func: (hass: HomeAssistant, entityId: string, ...args: any[]) => Promise<T>,
+  func: (hass: HomeAssistant, ...args: any[]) => Promise<T>,
+  generateCacheKey:
+    | ((hass: HomeAssistant, lastResult: T) => unknown)
+    | undefined,
   hass: HomeAssistant,
-  entityId: string,
   ...args: any[]
 ): Promise<T> => {
-  let cache: ResultCache<T> | undefined = (hass as any)[cacheKey];
+  const anyHass = hass as any;
+  const lastResult: Promise<CacheResult<T>> | CacheResult<T> | undefined =
+    anyHass[cacheKey];
 
-  if (!cache) {
-    cache = hass[cacheKey] = {};
-  }
+  const checkCachedResult = (result: CacheResult<T>): T | Promise<T> => {
+    if (
+      !generateCacheKey ||
+      generateCacheKey(hass, result.result) === result.cacheKey
+    ) {
+      return result.result;
+    }
 
-  const lastResult = cache[entityId];
+    anyHass[cacheKey] = undefined;
+    return timeCachePromiseFunc(
+      cacheKey,
+      cacheTime,
+      func,
+      generateCacheKey,
+      hass,
+      ...args
+    );
+  };
 
+  // If we have a cached result, return it if it's still valid
   if (lastResult) {
-    return lastResult;
+    return lastResult instanceof Promise
+      ? lastResult.then(checkCachedResult)
+      : checkCachedResult(lastResult);
   }
 
-  const result = func(hass, entityId, ...args);
-  cache[entityId] = result;
+  const resultPromise = func(hass, ...args);
+  anyHass[cacheKey] = resultPromise;
 
-  result.then(
+  resultPromise.then(
     // When successful, set timer to clear cache
-    () =>
+    (result) => {
+      anyHass[cacheKey] = {
+        result,
+        cacheKey: generateCacheKey?.(hass, result),
+      };
       setTimeout(() => {
-        cache![entityId] = undefined;
-      }, cacheTime),
+        anyHass[cacheKey] = undefined;
+      }, cacheTime);
+    },
     // On failure, clear cache right away
     () => {
-      cache![entityId] = undefined;
+      anyHass[cacheKey] = undefined;
     }
   );
 
-  return result;
+  return resultPromise;
 };

src/components/ha-selector/ha-selector-entity.ts

@@ -1,21 +1,23 @@
-import { HassEntity, UnsubscribeFunc } from "home-assistant-js-websocket";
-import { html, LitElement } from "lit";
+import { HassEntity } from "home-assistant-js-websocket";
+import { html, LitElement, PropertyValues } from "lit";
 import { customElement, property, state } from "lit/decorators";
 import { computeStateDomain } from "../../common/entity/compute_state_domain";
-import { subscribeEntityRegistry } from "../../data/entity_registry";
+import {
+  EntitySources,
+  fetchEntitySourcesWithCache,
+} from "../../data/entity_sources";
 import { EntitySelector } from "../../data/selector";
-import { SubscribeMixin } from "../../mixins/subscribe-mixin";
 import { HomeAssistant } from "../../types";
 import "../entity/ha-entities-picker";
 import "../entity/ha-entity-picker";
 
 @customElement("ha-selector-entity")
-export class HaEntitySelector extends SubscribeMixin(LitElement) {
+export class HaEntitySelector extends LitElement {
   @property() public hass!: HomeAssistant;
 
   @property() public selector!: EntitySelector;
 
-  @state() private _entityPlaformLookup?: Record<string, string>;
+  @state() private _entitySources?: EntitySources;
 
   @property() public value?: any;
 
@@ -49,49 +51,47 @@ export class HaEntitySelector extends SubscribeMixin(LitElement) {
     `;
   }
 
-  public hassSubscribe(): UnsubscribeFunc[] {
-    return [
-      subscribeEntityRegistry(this.hass.connection!, (entities) => {
-        const entityLookup = {};
-        for (const confEnt of entities) {
-          if (!confEnt.platform) {
-            continue;
-          }
-          entityLookup[confEnt.entity_id] = confEnt.platform;
-        }
-        this._entityPlaformLookup = entityLookup;
-      }),
-    ];
+  protected updated(changedProps: PropertyValues): void {
+    super.updated(changedProps);
+    if (
+      changedProps.has("selector") &&
+      this.selector.entity.integration &&
+      !this._entitySources
+    ) {
+      fetchEntitySourcesWithCache(this.hass).then((sources) => {
+        this._entitySources = sources;
+      });
+    }
   }
 
   private _filterEntities = (entity: HassEntity): boolean => {
-    if (this.selector.entity?.domain) {
-      const filterDomain = this.selector.entity.domain;
-      const filterDomainIsArray = Array.isArray(filterDomain);
+    const {
+      domain: filterDomain,
+      device_class: filterDeviceClass,
+      integration: filterIntegration,
+    } = this.selector.entity;
+
+    if (filterDomain) {
       const entityDomain = computeStateDomain(entity);
       if (
-        (filterDomainIsArray && !filterDomain.includes(entityDomain)) ||
-        (!filterDomainIsArray && entityDomain !== filterDomain)
+        Array.isArray(filterDomain)
+          ? !filterDomain.includes(entityDomain)
+          : entityDomain !== filterDomain
       ) {
         return false;
       }
     }
-    if (this.selector.entity?.device_class) {
-      if (
-        !entity.attributes.device_class ||
-        entity.attributes.device_class !== this.selector.entity.device_class
-      ) {
-        return false;
-      }
+    if (
+      filterDeviceClass &&
+      entity.attributes.device_class !== filterDeviceClass
+    ) {
+      return false;
     }
-    if (this.selector.entity?.integration) {
-      if (
-        !this._entityPlaformLookup ||
-        this._entityPlaformLookup[entity.entity_id] !==
-          this.selector.entity.integration
-      ) {
-        return false;
-      }
+    if (
+      filterIntegration &&
+      this._entitySources?.[entity.entity_id]?.domain !== filterIntegration
+    ) {
+      return false;
     }
     return true;
   };

src/data/camera.ts

@@ -2,7 +2,7 @@ import {
   HassEntityAttributeBase,
   HassEntityBase,
 } from "home-assistant-js-websocket";
-import { timeCachePromiseFunc } from "../common/util/time-cache-function-promise";
+import { timeCacheEntityPromiseFunc } from "../common/util/time-cache-entity-promise-func";
 import { HomeAssistant } from "../types";
 import { getSignedPath } from "./auth";
 
@@ -50,7 +50,7 @@ export const fetchThumbnailUrlWithCache = async (
   width: number,
   height: number
 ) => {
-  const base_url = await timeCachePromiseFunc(
+  const base_url = await timeCacheEntityPromiseFunc(
     "_cameraTmbUrl",
     9000,
     fetchThumbnailUrl,

src/data/entity_sources.ts

@@ -0,0 +1,46 @@
+import { timeCachePromiseFunc } from "../common/util/time-cache-function-promise";
+import { HomeAssistant } from "../types";
+
+interface EntitySourceConfigEntry {
+  source: "config_entry";
+  domain: string;
+  custom_component: boolean;
+  config_entry: string;
+}
+
+interface EntitySourcePlatformConfig {
+  source: "platform_config";
+  domain: string;
+  custom_component: boolean;
+}
+
+export type EntitySources = Record<
+  string,
+  EntitySourceConfigEntry | EntitySourcePlatformConfig
+>;
+
+const fetchEntitySources = (
+  hass: HomeAssistant,
+  entity_id?: string
+): Promise<EntitySources> =>
+  hass.callWS({
+    type: "entity/source",
+    entity_id,
+  });
+
+export const fetchEntitySourcesWithCache = (
+  hass: HomeAssistant,
+  entity_id?: string
+): Promise<EntitySources> =>
+  entity_id
+    ? fetchEntitySources(hass, entity_id)
+    : timeCachePromiseFunc(
+        "_entitySources",
+        // cache for 30 seconds
+        30000,
+        fetchEntitySources,
+        // We base the cache on number of states. If number of states
+        // changes we force a refresh
+        (hass2) => Object.keys(hass2.states).length,
+        hass
+      );