Signal store - How to get the types right for a custom store feature that uses factory-method?

[PhilippMDoerner]

Heyho, thought I'd make a new question as this is the result of this question.

I'm now trying to make a withQueries signal-store-feature that takes a factory function same as withMethods.
There the user defines a dictionary of queries - functions that return observables - whose data gets transferred to the store etc.

The fields I create are based on the keys on the returned object by the factory, so the syntax I'm going for is something like this:

export const SomeStore = signalStore(
  withState({ someVal: 'init' }),
  withQueries(store => {
    const service1 = inject(Service1);
    const service2 = inject(Service2);
    return {
      data1: (params: Params1) => service1(params),
      data2: (params: Params2) => service2(params, store.someVal()),
    } 
  });
)

Which should create a store with the added state signals:
someVal, data1, data1Errors, data1QueryState, data2, data2Errors, data2QueryState as well as the rxMethods loadData1 and loadData2.

Now I'm struggling to create a way to add that state from said factory function, as apparently the type of the store does not match what I'm expecting (?).

This is as far as I've gotten so far (Note, there are some import errors in the stackblitz that do not occur locally on my machine, notably when importing from '@ngrx/signals/src/signal-store-models'):

type QueriesFactory<
  Input extends SignalStoreFeatureResult,
  Methods extends MethodsDictionary,
> = (
  store: Prettify<
    StateSignals<Input['state']> &
      Input['computed'] &
      Input['methods'] &
      WritableStateSource<Prettify<Input['state']>>
  >,
) => Methods;

function withQueriesState<
  Input extends SignalStoreFeatureResult,
  InputMethods extends MethodsDictionary,
  State extends object,
>(
  queriesFactory: QueriesFactory<Input, InputMethods>,
): SignalStoreFeature<Input, { state: State; computed: {}; methods: {} }> {
  return ((store) => {
    const newState = queriesFactory({
      [STATE_SOURCE]: store[STATE_SOURCE],
      ...store.stateSignals,
      ...store.computedSignals,
      ...store.methods,
    });
    const queryStates = Object.keys(newState)
      .map((queryName) => getKeys(queryName))
      .map((keys) => ({
        [keys.dataField]: undefined,
        [keys.errorField]: undefined,
        [keys.queryStateField]: 'init' satisfies QueryState,
      }));

    const queryStateSignals = queryStates.reduce(
      (acc, queryState) => ({ ...acc, ...queryState }),
      {},
    );
    return withState(queryStateSignals)(store);
  }) as SignalStoreFeature<Input, { state: State; computed: {}; methods: {} }>;
}

export function withQueries<
  Input extends SignalStoreFeatureResult,
  State extends object,
  InputMethods extends MethodsDictionary,
  OutputMethods extends MethodsDictionary,
>(
  queriesFactory: QueriesFactory<Input, InputMethods>,
): SignalStoreFeature<
  Input,
  { state: State; computed: {}; methods: OutputMethods }
> {
  const stateFeature = withQueriesState(queriesFactory);

  return signalStoreFeature(
    (store) => stateFeature(store), // This causes type-errors because the type of store does not match what `queriesFactory` expects apparently
    withMethods((store) => { // Will create the load-methods once I've figured out how to add the state properly
      return {};
    }),
  );
}

[rainerhahnekamp]

@PhilippMDoerner did you try to use withMethods directly in withQueries?

[PhilippMDoerner]

I assume you mean withQueriesState ? withMethods is just exemplary here in that this is a placeholder for where I'll eventually add the load methods, the sole problem right now is withQueriesState.

It pretty much does not matter where I move the invocation of either of those, be that inside of the signalStoreFeatures call or outside, the type error in the feature created by withQueriesState remains the same.

export function withQueries<
  Input extends SignalStoreFeatureResult,
  State extends object,
  InputMethods extends MethodsDictionary,
  OutputMethods extends MethodsDictionary,
>(
  queriesFactory: QueriesFactory<Input, InputMethods>,
): SignalStoreFeature<
  Input,
  { state: State; computed: {}; methods: OutputMethods }
> {
  return signalStoreFeature(
   store =>  withQueriesState(queriesFactory)(store) //Same type error,
    withMethods((store) => {
      return {};
    }),
  );
}

// alternatively
export function withQueries
...
{
  const stateFeature = withQueriesState(queriesFactory);
  const methodsFeature = withMethods((store) => {
    return {};
  });
  return signalStoreFeature(
    withQueriesState(queriesFactory),
    methodsFeature,
  );
}

[PhilippMDoerner]

Trying to get this to work is proving trickier than I thought. Out of curiosity, are there any other custom feature hooks that you're aware of that I could take a look at that function similarly to withMethods in that they use a factory-function ala store => object with properties ?

I keep getting the types wrong.

Edit: Maybe to elaborate: The complexity that stumps me is the fact that unlike withMethods I don't just want to transfer key-value pairs into a store, I want to take the keys, transform those keys into different keys and use those for property names.

I can do that as long as I have individual elements for which I want to generate fields (1 query).
Once the user passes in multiple inputs (an object with multiple queries or even a factory function) I'm getting somewhat stumped on how the types need to look in order to work and be correct.

Edit: This is as far as I've gotten anyway: https://tinyurl.com/mwrk2nbn

[PhilippMDoerner]

I've made... some progress and have a first functional draft of a feature, but I'm really starting to feel like I'm scratching at even the limits of typescript itself (see this typescript discord post and this discussion within the angular community discord ).
I'm certainly out of my depth in that regard, twice over.

Managed to make it work though by creating the state/method object inside a withState/withMethods-callbacks without specific types, then type-casting those to a generic type (AllNewProperties / AllNewMethods ) that adds all the fields to the state/method-type of the store as well.

Given the sheer amount of code solely for types, I split things up into 4 files:

• index.ts for the final feature
• withQueryState to add the state-fields
• withQueryMethods to add the load-methods
• types.ts for the getKeys function (that both withQueryState and withQueryMethods need) as well as some shared types ( QueryState, Query, QueryMap etc.).

/withQueries/withQueryState.ts

import { HttpErrorResponse } from '@angular/common/http';
import { signalStoreFeature, withState } from '@ngrx/signals';
import {
  getKeys,
  Query,
  QueryMap,
  QueryState,
  SomeVersionOfU2I,
} from './types';

// Creates an object with a bunch of properties based on an input name
type NewProperties<Name extends string, Q> =
  Q extends Query<infer Params, infer Response>
    ? Record<Uncapitalize<Name>, Response | undefined> &
        Record<`${Name}Error`, HttpErrorResponse> &
        Record<`${Uncapitalize<Name>}QueryState`, QueryState>
    : never;

// I think represents a single slice of:
// `{ <queryName>: { <queyName>: Response, <queryName>Error: HttpErrorResponse, <queryName>QueryState: QueryState } }`
type SingleNewPropertiesObjectSlice<Queries extends QueryMap<unknown>> = {
  [Key in keyof Queries & string]: NewProperties<
    Key,
    SomeVersionOfU2I<Queries[Key]>
  >;
};

// Unpacked version of `SingleNewPropertiesObjectSlice`, essentially just `{ <queyName>: Response, <queryName>Error: HttpErrorResponse, <queryName>QueryState: QueryState }`
type SingleNewProperty<Queries extends QueryMap<unknown>> =
  SingleNewPropertiesObjectSlice<Queries>[keyof SingleNewPropertiesObjectSlice<Queries>];

// Extracts params from queries and that... interacts.. somehow to make the correct type
export type AllNewProperties<Queries extends QueryMap<unknown>> =
  SomeVersionOfU2I<SingleNewProperty<Queries>>;

export function withQueriesState<Queries extends QueryMap<any>>(
  queriesFactory: () => Queries,
) {
  return signalStoreFeature(
    withState(() => {
      const queries = queriesFactory();
      const queryStates = Object.keys(queries)
        .map((queryName) => getKeys(queryName))
        .map((keys) => {
          const x = {
            [keys.dataField]: undefined,
            [keys.errorField]: undefined,
            [keys.queryStateField]: 'init' satisfies QueryState,
          } as SingleNewProperty<Queries>;
          return x;
        });

      const queryStateSignals = queryStates.reduce(
        (acc, queryState) => ({ ...acc, ...queryState }),
        {},
      ) as AllNewProperties<Queries> & {};
      return queryStateSignals;
    }),
  );
}

/withQueries/withQueryMethods.ts

import { tapResponse } from '@ngrx/operators';
import { patchState, signalStoreFeature, withMethods } from '@ngrx/signals';
import { rxMethod } from '@ngrx/signals/rxjs-interop';
import { MethodsDictionary } from '@ngrx/signals/src/signal-store-models';
import { pipe, switchMap, tap } from 'rxjs';
import {
  getKeys,
  Query,
  QueryMap,
  QueryState,
  SomeVersionOfU2I,
} from './types';

type NewMethods<Name extends string, Q> =
  Q extends Query<infer Params, infer Response>
    ? Record<`load${Capitalize<Name>}`, ReturnType<typeof rxMethod<Params>>>
    : never;

type SingleNewMethodObjectSlice<Queries extends QueryMap<unknown>> = {
  [Key in keyof Queries & string]: NewMethods<
    Key,
    SomeVersionOfU2I<Queries[Key]>
  >;
};

type SingleNewMethod<Queries extends QueryMap<unknown>> =
  SingleNewMethodObjectSlice<Queries>[keyof SingleNewMethodObjectSlice<Queries>];

// Extracts params from queries and that... interacts.. somehow to make the correct type
export type AllNewMethods<Queries extends QueryMap<unknown>> = SomeVersionOfU2I<
  SingleNewMethod<Queries>
>;

export function withQueryMethods<Queries extends QueryMap<any>>(
  queriesFactory: () => Queries,
) {
  return signalStoreFeature(
    withMethods((store) => {
      const queries = queriesFactory();
      const queryKeys = Object.keys(queries).map((queryName) =>
        getKeys(queryName),
      );

      const queryLoadFunctions = queryKeys.map((keys) => {
        return {
          [keys.loadMethod]: rxMethod(
            pipe(
              tap(() =>
                patchState(store, {
                  [keys.queryStateField]: 'loading' satisfies QueryState,
                }),
              ),
              switchMap((params) => queries[keys.name](params)),
              tapResponse({
                next: (val) =>
                  patchState(store, {
                    [keys.dataField]: val,
                    [keys.queryStateField]: 'success' satisfies QueryState,
                  } as any),
                error: (err) =>
                  patchState(store, {
                    [keys.errorField]: err,
                    [keys.queryStateField]: 'error' satisfies QueryState,
                  } as any),
              }),
            ),
          ),
        };
      });

      const functions = queryLoadFunctions.reduce(
        (acc, queryLoadFunction) => ({ ...acc, ...queryLoadFunction }),
        {},
      ) as MethodsDictionary & AllNewMethods<Queries>;
      return functions;
    }),
  );
}

/withQueries/types.ts

import { Observable } from 'rxjs';
import { capitalize, uncapitalize } from 'src/utils/string';

export type SomeVersionOfU2I<U> = (
  U extends any ? (x: U) => any : never
) extends (x: infer I) => any
  ? I
  : never;

export type QueryState = 'init' | 'loading' | 'success' | 'error';
export type Query<Params, Response> = (params: Params) => Observable<Response>;
export type QueryMap<T extends any | unknown> = Record<string, Query<T, T>>;

export function getKeys<Name extends string>(
  name: Name,
): Record<'name', Name> &
  Record<'dataField', Uncapitalize<Name>> &
  Record<'errorField', `${Uncapitalize<Name>}Error`> &
  Record<'queryStateField', `${Uncapitalize<Name>}QueryState`> &
  Record<'loadMethod', `load${Capitalize<Name>}`> {
  return {
    name,
    dataField: uncapitalize(name),
    errorField: `${uncapitalize(name)}Error`,
    queryStateField: `${uncapitalize(name)}QueryState`,
    loadMethod: `load${capitalize(name)}`,
  };
}

/withQueries/index.ts

import { signalStoreFeature } from '@ngrx/signals';
import { QueryMap } from './types';
import { withQueryMethods } from './withQueryMethods';
import { withQueriesState } from './withQueryState';

// The types below are useless as `AllNewMethods` and `AllNewProperties` inside them get evaluated to unknown, but work inside the 2 sub-features
// type QueriesFeatureResult<Queries extends QueryMap<any>> = {
//   computed: {};
//   methods: MethodsDictionary & AllNewMethods<Queries>;
//   state: {} & AllNewProperties<Queries>;
// };

// type QueriesFeature<Queries extends QueryMap<any>> = SignalStoreFeature<
//   EmptyFeatureResult,
//   QueriesFeatureResult<Queries>
// >;

export function withQueries<Queries extends QueryMap<any>>(
  queriesFactory: () => Queries,
) {
  return signalStoreFeature(
    withQueriesState(queriesFactory),
    withQueryMethods(queriesFactory),
  );
}

This nets you a feature you can use the same as withMethods and similarly to withMethods you get to do DI inside that function.

Example Usage:

export const MyStore = signalStore(
  withQueries(() => {
    const charService = inject(CharacterService);
    return {
      characters: (params: { name: string }) => {
        return charService.readByParam(params);
      },
    };
  }),
)

Limitations of this approach that I can't seem to find a way around: the withQueries callback can't access store, because there is no version of withState that takes in a callback-function that provides a store. So if you want to trigger something else inside the store upon calling a query function or when it gets resolved, then you can't. At least not directly inside `withQueries.

I'll post an update if I ever figure out how/if that can be simplified. This certainly showed me that what I'm trying to do here appears to be anything but simple.

[PhilippMDoerner]

After extensive discussions about the depths of Typescript and the approaches you can use in various scenarios within the Angular discord , the consensus was pretty much: I must use type-casting to establish which properties get added to the store, as TS can not properly infer those for me. The key problem for that inference is Object.keys(queries) as that turns all the property names into the type "string[]" instead of a type-union of property-names.

Since that is the case, I kind of just started brute-forcing the issue, which worked a lot easier than anticipated:

/withQueries/index.ts

import { WritableStateSource } from '@ngrx/signals';
import {
  MethodsDictionary,
  SignalStoreFeature,
  SignalStoreFeatureResult,
  StateSignals,
} from '@ngrx/signals/src/signal-store-models';
import { QueryMap } from './types';
import { AllNewMethods, withQueryMethods } from './withQueryMethods';
import { AllNewProperties, withQueriesState } from './withQueryState';

type QueriesFeatureResult<Queries extends QueryMap> = {
  computed: {};
  methods: MethodsDictionary & AllNewMethods<Queries>;
  state: {} & AllNewProperties<Queries>;
};

type InnerStore<Input extends SignalStoreFeatureResult> = StateSignals<
  Input['state']
> &
  Input['computed'] &
  Input['methods'] &
  WritableStateSource<Input['state']>;

export function withQueries<
  Input extends SignalStoreFeatureResult,
  Queries extends QueryMap,
>(
  queriesFactory: (store: InnerStore<Input>) => Queries,
): SignalStoreFeature<Input, QueriesFeatureResult<Queries>> {
  return ((store) => {
    const queries = queriesFactory({
      ...store.stateSignals,
      ...store.computedSignals,
      ...store.methods,
    } as InnerStore<Input>);

    const storeWithState = withQueriesState(queries)(store);
    return withQueryMethods(queries)(storeWithState);
  }) as SignalStoreFeature<Input, QueriesFeatureResult<Queries>>;
}

/withQueries/withQueryState.ts

import { HttpErrorResponse } from '@angular/common/http';
import { signalStoreFeature, withState } from '@ngrx/signals';
import {
  getKeys,
  Query,
  QueryMap,
  QueryState,
  SomeVersionOfU2I,
} from './types';

// Creates an object with a bunch of properties based on an input name
type NewProperties<Name extends string, Q> =
  Q extends Query<infer Params, infer Response>
    ? Record<Uncapitalize<Name>, Response | undefined> &
        Record<`${Name}Error`, HttpErrorResponse> &
        Record<`${Uncapitalize<Name>}QueryState`, QueryState>
    : never;

// I think represents a single slice of:
// `{ <queryName>: { <queyName>: Response, <queryName>Error: HttpErrorResponse, <queryName>QueryState: QueryState } }`
type SingleNewPropertiesObjectSlice<Queries extends QueryMap> = {
  [Key in keyof Queries & string]: NewProperties<
    Key,
    SomeVersionOfU2I<Queries[Key]>
  >;
};

// Unpacked version of `SingleNewPropertiesObjectSlice`, essentially just `{ <queyName>: Response, <queryName>Error: HttpErrorResponse, <queryName>QueryState: QueryState }`
type SingleNewProperty<Queries extends QueryMap> =
  SingleNewPropertiesObjectSlice<Queries>[keyof SingleNewPropertiesObjectSlice<Queries>];

// Extracts params from queries and that... interacts.. somehow to make the correct type
export type AllNewProperties<Queries extends QueryMap> = SomeVersionOfU2I<
  SingleNewProperty<Queries>
>;

export function withQueriesState<Queries extends QueryMap>(queries: Queries) {
  return signalStoreFeature(
    withState(() => {
      const queryStates = Object.keys(queries)
        .map((queryName) => getKeys(queryName))
        .map((keys) => {
          const x = {
            [keys.dataField]: undefined,
            [keys.errorField]: undefined,
            [keys.queryStateField]: 'init' satisfies QueryState,
          };
          return x;
        });

      const queryStateSignals = queryStates.reduce(
        (acc, queryState) => ({ ...acc, ...queryState }),
        {},
      ) as AllNewProperties<Queries> & {};
      return queryStateSignals;
    }),
  );
}

/withQueries/withQueryMethods.ts

import { tapResponse } from '@ngrx/operators';
import { patchState, signalStoreFeature, withMethods } from '@ngrx/signals';
import { rxMethod } from '@ngrx/signals/rxjs-interop';
import { MethodsDictionary } from '@ngrx/signals/src/signal-store-models';
import { pipe, switchMap, take, tap } from 'rxjs';
import {
  getKeys,
  Query,
  QueryMap,
  QueryState,
  SomeVersionOfU2I,
} from './types';

type NewMethods<Name extends string, Q> =
  Q extends Query<void, any>
    ? Record<`load${Capitalize<Name>}`, ReturnType<typeof rxMethod<void>>>
    : Q extends Query<infer Params, any>
      ? Record<`load${Capitalize<Name>}`, ReturnType<typeof rxMethod<Params>>>
      : never;

type SingleNewMethodObjectSlice<Queries extends QueryMap> = {
  [Key in keyof Queries & string]: NewMethods<
    Key,
    SomeVersionOfU2I<Queries[Key]>
  >;
};

type SingleNewMethod<Queries extends QueryMap> =
  SingleNewMethodObjectSlice<Queries>[keyof SingleNewMethodObjectSlice<Queries>];

// Extracts params from queries and that... interacts.. somehow to make the correct type
export type AllNewMethods<Queries extends QueryMap> = SomeVersionOfU2I<
  SingleNewMethod<Queries>
>;

export function withQueryMethods<Queries extends QueryMap>(queries: Queries) {
  return signalStoreFeature(
    withMethods((store) => {
      const queryLoadFunctions = Object.keys(queries)
        .map((queryName) => getKeys(queryName))
        .map((keys) => {
          return {
            [keys.loadMethod]: rxMethod(
              pipe(
                tap(() =>
                  patchState(store, {
                    [keys.queryStateField]: 'loading' satisfies QueryState,
                  }),
                ),
                switchMap((params) => queries[keys.name](params)),
                take(1),
                tapResponse({
                  next: (val) =>
                    patchState(store, {
                      [keys.dataField]: val,
                      [keys.queryStateField]: 'success' satisfies QueryState,
                    }),
                  error: (err) =>
                    patchState(store, {
                      [keys.errorField]: err,
                      [keys.queryStateField]: 'error' satisfies QueryState,
                    }),
                }),
              ),
            ),
          };
        });

      const functions = queryLoadFunctions.reduce(
        (acc, queryLoadFunction) => ({ ...acc, ...queryLoadFunction }),
        {},
      ) as MethodsDictionary & AllNewMethods<Queries>;
      return functions;
    }),
  );
}

/withQueries/types.ts

import { Observable } from 'rxjs';
import { capitalize, uncapitalize } from 'src/utils/string';

export type SomeVersionOfU2I<U> = (
  U extends any ? (x: U) => any : never
) extends (x: infer I) => any
  ? I
  : never;

export type QueryState = 'init' | 'loading' | 'success' | 'error';
export type Query<Params, Response> = (params: Params) => Observable<Response>;
export type QueryMap = Record<string, Query<any, any>>;

export function getKeys<Name extends string>(
  name: Name,
): Record<'name', Name> &
  Record<'dataField', Uncapitalize<Name>> &
  Record<'errorField', `${Uncapitalize<Name>}Error`> &
  Record<'queryStateField', `${Uncapitalize<Name>}QueryState`> &
  Record<'loadMethod', `load${Capitalize<Name>}`> {
  return {
    name,
    dataField: uncapitalize(name),
    errorField: `${uncapitalize(name)}Error`,
    queryStateField: `${uncapitalize(name)}QueryState`,
    loadMethod: `load${capitalize(name)}`,
  };
}