feat(signals): make EntityId type an optional argument in withEntities

modules/signals/entities/spec/types/with-entities.types.spec.ts

@@ -0,0 +1,173 @@
+import { expecter } from 'ts-snippet';
+import { compilerOptions } from './helpers';
+
+describe('withEntities', () => {
+  const expectSnippet = expecter(
+    (code) => `
+      import { signalStore, type } from '@ngrx/signals';
+      import { withEntities, entityConfig} from '@ngrx/signals/entities';
+
+      enum UserId {
+        One = '1',
+        Two = '2',
+        Three = '3',
+      }
+
+      ${code}
+    `,
+    compilerOptions()
+  );
+
+  it('succeeds when only Entity type is provided as a generic', () => {
+    const snippet = `
+      type User = { id: UserId; name: string };
+
+      const Store = signalStore(
+        withEntities<User>()
+      );
+
+      const store = new Store();
+      const ids = store.ids();
+      const entityMap = store.entityMap();
+    `;
+
+    expectSnippet(snippet).toSucceed();
+    expectSnippet(snippet).toInfer('ids', 'EntityId[]');
+    expectSnippet(snippet).toInfer('entityMap', 'EntityMap<User, EntityId>');
+  });
+
+  it('succeeds when both Entity and Id types are provided as generics', () => {
+    const snippet = `
+      type User = { id: UserId; name: string };
+
+      const Store = signalStore(
+        withEntities<User, UserId>()
+      );
+
+      const store = new Store();
+      const ids = store.ids();
+      const entityMap = store.entityMap();
+    `;
+
+    expectSnippet(snippet).toSucceed();
+    expectSnippet(snippet).toInfer('ids', 'UserId[]');
+    expectSnippet(snippet).toInfer('entityMap', 'EntityMap<User, UserId>');
+  });
+
+  it('succeeds when Entity type is provided using entityConfig', () => {
+    const snippet = `
+      type User = { id: UserId; name: string };
+
+      const Store = signalStore(
+        withEntities(entityConfig({entity: type<User>()}))
+      );
+
+      const store = new Store();
+      const ids = store.ids();
+      const entityMap = store.entityMap();
+    `;
+
+    expectSnippet(snippet).toSucceed();
+    expectSnippet(snippet).toInfer('ids', 'EntityId[]');
+    expectSnippet(snippet).toInfer('entityMap', 'EntityMap<User, EntityId>');
+  });
+
+  it('succeeds when Entity type is provided as both a generic and in entityConfig', () => {
+    const snippet = `
+      type User = { id: UserId; name: string };
+
+      const Store = signalStore(
+        withEntities<User>(entityConfig({entity: type<User>()}))
+      );
+
+      const store = new Store();
+      const ids = store.ids();
+      const entityMap = store.entityMap();
+    `;
+
+    expectSnippet(snippet).toSucceed();
+    expectSnippet(snippet).toInfer('ids', 'EntityId[]');
+    expectSnippet(snippet).toInfer('entityMap', 'EntityMap<User, EntityId>');
+  });
+
+  it('succeeds when Entity type is provided as both a generic and in entityConfig, with Id type as a generic', () => {
+    const snippet = `
+      type User = { id: UserId; name: string };
+
+      const Store = signalStore(
+        withEntities<User, UserId>(entityConfig({entity: type<User>()}))
+      );
+
+      const store = new Store();
+      const ids = store.ids();
+      const entityMap = store.entityMap();
+    `;
+
+    expectSnippet(snippet).toSucceed();
+    expectSnippet(snippet).toInfer('ids', 'UserId[]');
+    expectSnippet(snippet).toInfer('entityMap', 'EntityMap<User, UserId>');
+  });
+
+  it('succeeds when Entity type is provided with a custom selectId function', () => {
+    const snippet = `
+      type User = { key: UserId; name: string };
+
+      const Store = signalStore(
+        withEntities<User, UserId>(entityConfig({
+          entity: type<User>(),
+          selectId: (user) => user.key
+        }))
+      );
+
+      const store = new Store();
+      const ids = store.ids();
+      const entityMap = store.entityMap();
+    `;
+
+    expectSnippet(snippet).toSucceed();
+    expectSnippet(snippet).toInfer('ids', 'UserId[]');
+    expectSnippet(snippet).toInfer('entityMap', 'EntityMap<User, UserId>');
+  });
+
+  it('succeeds when Entity type is provided with a custom collection name', () => {
+    const snippet = `
+      type User = { id: UserId; name: string };
+
+      const Store = signalStore(
+        withEntities(entityConfig({
+          entity: type<User>(),
+          collection: 'user'
+        }))
+      );
+
+      const store = new Store();
+      const ids = store.userIds();
+      const entityMap = store.userEntityMap();
+    `;
+
+    expectSnippet(snippet).toSucceed();
+    expectSnippet(snippet).toInfer('ids', 'EntityId[]');
+    expectSnippet(snippet).toInfer('entityMap', 'EntityMap<User, EntityId>');
+  });
+
+  it('succeeds when Entity, Id type, and custom collection name are provided', () => {
+    const snippet = `
+      type User = { id: UserId; name: string };
+
+      const Store = signalStore(
+        withEntities<User, 'user', UserId>(entityConfig({
+          entity: type<User>(),
+          collection: 'user'
+        }))
+      );
+
+      const store = new Store();
+      const ids = store.userIds();
+      const entityMap = store.userEntityMap();
+    `;
+
+    expectSnippet(snippet).toSucceed();
+    expectSnippet(snippet).toInfer('ids', 'UserId[]');
+    expectSnippet(snippet).toInfer('entityMap', 'EntityMap<User, UserId>');
+  });
+});

modules/signals/entities/src/models.ts

@@ -2,15 +2,25 @@ import { Signal } from '@angular/core';
 
 export type EntityId = string | number;
 
-export type EntityMap<Entity> = Record<EntityId, Entity>;
-
-export type EntityState<Entity> = {
-  entityMap: EntityMap<Entity>;
-  ids: EntityId[];
+export type EntityMap<Entity, Id extends EntityId = EntityId> = Record<
+  Id,
+  Entity
+>;
+
+export type EntityState<Entity, Id extends EntityId = EntityId> = {
+  entityMap: EntityMap<Entity, Id>;
+  ids: Id[];
 };
 
-export type NamedEntityState<Entity, Collection extends string> = {
-  [K in keyof EntityState<Entity> as `${Collection}${Capitalize<K>}`]: EntityState<Entity>[K];
+export type NamedEntityState<
+  Entity,
+  Collection extends string,
+  Id extends EntityId = EntityId
+> = {
+  [K in keyof EntityState<
+    Entity,
+    Id
+  > as `${Collection}${Capitalize<K>}`]: EntityState<Entity, Id>[K];
 };
 
 export type EntityProps<Entity> = {

modules/signals/entities/src/with-entities.ts

@@ -16,36 +16,43 @@ import {
 } from './models';
 import { getEntityStateKeys } from './helpers';
 
-export function withEntities<Entity>(): SignalStoreFeature<
+export function withEntities<
+  Entity,
+  Id extends EntityId = EntityId
+>(): SignalStoreFeature<
   EmptyFeatureResult,
   {
-    state: EntityState<Entity>;
+    state: EntityState<Entity, Id>;
     props: EntityProps<Entity>;
     methods: {};
   }
 >;
-export function withEntities<Entity, Collection extends string>(config: {
+export function withEntities<
+  Entity,
+  Collection extends string,
+  Id extends EntityId = EntityId
+>(config: {
   entity: Entity;
   collection: Collection;
 }): SignalStoreFeature<
   EmptyFeatureResult,
   {
-    state: NamedEntityState<Entity, Collection>;
+    state: NamedEntityState<Entity, Collection, Id>;
     props: NamedEntityProps<Entity, Collection>;
     methods: {};
   }
 >;
-export function withEntities<Entity>(config: {
+export function withEntities<Entity, Id extends EntityId = EntityId>(config: {
   entity: Entity;
 }): SignalStoreFeature<
   EmptyFeatureResult,
   {
-    state: EntityState<Entity>;
+    state: EntityState<Entity, Id>;
     props: EntityProps<Entity>;
     methods: {};
   }
 >;
-export function withEntities<Entity>(config?: {
+export function withEntities<Entity, Id extends EntityId = EntityId>(config?: {
   entity: Entity;
   collection?: string;
 }): SignalStoreFeature {
@@ -58,8 +65,8 @@ export function withEntities<Entity>(config?: {
     }),
     withComputed((store: Record<string, Signal<unknown>>) => ({
       [entitiesKey]: computed(() => {
-        const entityMap = store[entityMapKey]() as EntityMap<Entity>;
-        const ids = store[idsKey]() as EntityId[];
+        const entityMap = store[entityMapKey]() as EntityMap<Entity, Id>;
+        const ids = store[idsKey]() as Id[];
 
         return ids.map((id) => entityMap[id]);
       }),

projects/ngrx.io/content/guide/signals/signal-store/entity-management.md

@@ -34,6 +34,12 @@ The `withEntities` feature adds the following signals to the `TodosStore`:
 
 The `ids` and `entityMap` are state slices, while `entities` is a computed signal.
 
+If you need to use a custom type for `ids` and `entityMap` keys that extends `number` or `string`, you can pass this type as a generic:
+
+```ts
+withEntities<User, UserId>();
+```
+
 ## Entity Updaters
 
 The `entities` plugin provides a set of standalone entity updaters.