Merge pull request #986 from dermotduffy/no-longer-need-extended-for-uniqueid

Remove unnecessary entity fetches

src/utils/ha/entity-registry/cache.ts

@@ -1,7 +1,7 @@
-import { Entity, ExtendedEntity } from './types.js';
+import { Entity } from './types.js';
 
-export class EntityCache<T extends Entity | ExtendedEntity> {
-  protected _cache: Map<string, T> = new Map();
+export class EntityCache {
+  protected _cache: Map<string, Entity> = new Map();
 
   /**
    * Determine if the cache has a given entity_id.
@@ -21,25 +21,25 @@ export class EntityCache<T extends Entity | ExtendedEntity> {
   //   return [...this._cache.values()].find(func) ?? null;
   // }
 
-  public getMatches(func: (arg: T) => boolean): T[] {
+  public getMatches(func: (arg: Entity) => boolean): Entity[] {
     return [...this._cache.values()].filter(func);
   }
 
   /**
    * Get entity information given an id.
    * @param id The entity id.
-   * @returns The `ExtendedEntity` for this id.
+   * @returns The entity for this id.
    */
-  public get(id: string): T | undefined {
+  public get(id: string): Entity | undefined {
     return this._cache.get(id);
   }
 
   /**
    * Add a given entity to the cache.
-   * @param extendedEntity
+   * @param input The entity.
    */
-  public set(input: T | T[]): void {
-    const _set = (entity: T) => this._cache.set(entity.entity_id, entity);
+  public set(input: Entity | Entity[]): void {
+    const _set = (entity: Entity) => this._cache.set(entity.entity_id, entity);
 
     if (Array.isArray(input)) {
       input.forEach(_set);

src/utils/ha/entity-registry/index.ts

@@ -1,32 +1,18 @@
 import { HomeAssistant } from 'custom-card-helpers';
 import { homeAssistantWSRequest } from '..';
 import { EntityCache } from './cache';
-import {
-  Entity,
-  EntityList,
-  entityListSchema,
-  ExtendedEntity,
-  extendedEntitySchema,
-} from './types.js';
+import { Entity, EntityList, entitySchema, entityListSchema } from './types.js';
 
-type EntityRegistryCache = EntityCache<Entity>;
-type ExtendedEntityRegistryCache = EntityCache<ExtendedEntity>;
-
-// Tne `entity_registry/list` call returns a smaller set of information for
-// every entity, than the full `entity_registry/get` call returns for a single
-// entity. This class manages interactions with entities, caching results
-// (either the partial or extended versions) and fetching as necessary. Some
-// calls require every entity to be fetched, which may be non-trivial in size
-// (after which it is cached forever).
+// This class manages interactions with entities, caching results and fetching
+// as necessary. Some calls require every entity to be fetched, which may be
+// non-trivial in size (after which it is cached forever).
 
 export class EntityRegistryManager {
-  protected _cache: EntityRegistryCache;
-  protected _extendedCache: ExtendedEntityRegistryCache;
+  protected _cache: EntityCache;
   protected _fetchedEntityList = false;
 
-  constructor(cache: EntityCache<Entity>, extendedCache: EntityCache<ExtendedEntity>) {
+  constructor(cache: EntityCache) {
     this._cache = cache;
-    this._extendedCache = extendedCache;
   }
 
   public async getEntity(hass: HomeAssistant, entityID: string): Promise<Entity | null> {
@@ -35,11 +21,12 @@ export class EntityRegistryManager {
       return cachedEntity;
     }
 
-    const cachedExtendedEntity = this._extendedCache.get(entityID);
-    if (cachedExtendedEntity) {
-      return cachedExtendedEntity;
-    }
-    return await this.getExtendedEntity(hass, entityID);
+    const entity = await homeAssistantWSRequest<Entity>(hass, entitySchema, {
+      type: 'config/entity_registry/get',
+      entity_id: entityID,
+    });
+    this._cache.set(entity);
+    return entity;
   }
 
   public async getMatchingEntities(
@@ -50,26 +37,6 @@ export class EntityRegistryManager {
     return this._cache.getMatches(func);
   }
 
-  public async getExtendedEntity(
-    hass: HomeAssistant,
-    entityID: string,
-  ): Promise<ExtendedEntity> {
-    const cachedValue = this._extendedCache.get(entityID);
-    if (cachedValue) {
-      return cachedValue;
-    }
-    const extendedEntity = await homeAssistantWSRequest<ExtendedEntity>(
-      hass,
-      extendedEntitySchema,
-      {
-        type: 'config/entity_registry/get',
-        entity_id: entityID,
-      },
-    );
-    this._extendedCache.set(extendedEntity);
-    return extendedEntity;
-  }
-
   public async getEntities(
     hass: HomeAssistant,
     entityIDs: string[],
@@ -85,21 +52,6 @@ export class EntityRegistryManager {
     return output;
   }
 
-  public async getExtendedEntities(
-    hass: HomeAssistant,
-    entityIDs: string[],
-  ): Promise<Map<string, ExtendedEntity>> {
-    const output: Map<string, ExtendedEntity> = new Map();
-    const _storeExtendedEntity = async (entityID: string): Promise<void> => {
-      const extendedEntity = await this.getExtendedEntity(hass, entityID);
-      if (extendedEntity) {
-        output.set(entityID, extendedEntity);
-      }
-    };
-    await Promise.all(entityIDs.map(_storeExtendedEntity));
-    return output;
-  }
-
   public async fetchEntityList(hass: HomeAssistant): Promise<void> {
     if (this._fetchedEntityList) {
       return;

src/utils/ha/entity-registry/types.ts

@@ -1,19 +1,14 @@
 import { z } from 'zod';
 
-const entitySchema = z.object({
+export const entitySchema = z.object({
   config_entry_id: z.string().nullable(),
   disabled_by: z.string().nullable(),
   entity_id: z.string(),
   hidden_by: z.string().nullable(),
   platform: z.string(),
+  unique_id: z.string().optional(),
 });
 export type Entity = z.infer<typeof entitySchema>;
 
-export const extendedEntitySchema = entitySchema.extend({
-  // Extended entity results.
-  unique_id: z.string().optional(),
-});
-export type ExtendedEntity = z.infer<typeof extendedEntitySchema>;
-
 export const entityListSchema = entitySchema.array();
 export type EntityList = z.infer<typeof entityListSchema>;