add inference of raw result type from schema (or specifying from a third type parameter)

Author: EskiMojo14

packages/toolkit/src/query/core/apiState.ts

@@ -179,7 +179,7 @@ export type MutationKeys<Definitions extends EndpointDefinitions> = {
 }[keyof Definitions]
 
 type BaseQuerySubState<
-  D extends BaseEndpointDefinition<any, any, any>,
+  D extends BaseEndpointDefinition<any, any, any, any>,
   DataType = ResultTypeFrom<D>,
 > = {
   /**
@@ -217,7 +217,7 @@ type BaseQuerySubState<
 }
 
 export type QuerySubState<
-  D extends BaseEndpointDefinition<any, any, any>,
+  D extends BaseEndpointDefinition<any, any, any, any>,
   DataType = ResultTypeFrom<D>,
 > = Id<
   | ({ status: QueryStatus.fulfilled } & WithRequiredProp<
@@ -244,15 +244,17 @@ export type QuerySubState<
 export type InfiniteQueryDirection = 'forward' | 'backward'
 
 export type InfiniteQuerySubState<
-  D extends BaseEndpointDefinition<any, any, any>,
+  D extends BaseEndpointDefinition<any, any, any, any>,
 > =
   D extends InfiniteQueryDefinition<any, any, any, any, any>
     ? QuerySubState<D, InfiniteData<ResultTypeFrom<D>, PageParamFrom<D>>> & {
         direction?: InfiniteQueryDirection
       }
     : never
 
-type BaseMutationSubState<D extends BaseEndpointDefinition<any, any, any>> = {
+type BaseMutationSubState<
+  D extends BaseEndpointDefinition<any, any, any, any>,
+> = {
   requestId: string
   data?: ResultTypeFrom<D>
   error?:
@@ -265,8 +267,12 @@ type BaseMutationSubState<D extends BaseEndpointDefinition<any, any, any>> = {
   fulfilledTimeStamp?: number
 }
 
-export type MutationSubState<D extends BaseEndpointDefinition<any, any, any>> =
-  | (({ status: QueryStatus.fulfilled } & WithRequiredProp<
+export type MutationSubState<
+  D extends BaseEndpointDefinition<any, any, any, any>,
+> =
+  | (({
+      status: QueryStatus.fulfilled
+    } & WithRequiredProp<
       BaseMutationSubState<D>,
       'data' | 'fulfilledTimeStamp'
     >) & { error: undefined })

packages/toolkit/src/query/core/buildMiddleware/cacheLifecycle.ts

@@ -1,5 +1,9 @@
 import type { ThunkDispatch, UnknownAction } from '@reduxjs/toolkit'
-import type { BaseQueryFn, BaseQueryMeta } from '../../baseQueryTypes'
+import type {
+  BaseQueryFn,
+  BaseQueryMeta,
+  BaseQueryResult,
+} from '../../baseQueryTypes'
 import type { BaseEndpointDefinition } from '../../endpointDefinitions'
 import { DefinitionType, isAnyQueryDefinition } from '../../endpointDefinitions'
 import type { QueryCacheKey, RootState } from '../apiState'
@@ -32,7 +36,8 @@ export interface QueryBaseLifecycleApi<
     { type: DefinitionType.query } & BaseEndpointDefinition<
       QueryArg,
       BaseQuery,
-      ResultType
+      ResultType,
+      BaseQueryResult<BaseQuery>
     >
   >
   /**
@@ -55,7 +60,8 @@ export type MutationBaseLifecycleApi<
     { type: DefinitionType.mutation } & BaseEndpointDefinition<
       QueryArg,
       BaseQuery,
-      ResultType
+      ResultType,
+      BaseQueryResult<BaseQuery>
     >
   >
 }

packages/toolkit/src/query/endpointDefinitions.ts

@@ -46,6 +46,7 @@ type EndpointDefinitionWithQuery<
   QueryArg,
   BaseQuery extends BaseQueryFn,
   ResultType,
+  RawResultType extends BaseQueryResult<BaseQuery>,
 > = {
   /**
    * `query` can be a function that returns either a `string` or an `object` which is passed to your `baseQuery`. If you are using [fetchBaseQuery](./fetchBaseQuery), this can return either a `string` or an `object` of properties in `FetchArgs`. If you use your own custom [`baseQuery`](../../rtk-query/usage/customizing-queries), you can customize this behavior to your liking.
@@ -91,7 +92,7 @@ type EndpointDefinitionWithQuery<
    * A function to manipulate the data returned by a query or mutation.
    */
   transformResponse?(
-    baseQueryReturnValue: BaseQueryResult<BaseQuery>,
+    baseQueryReturnValue: RawResultType,
     meta: BaseQueryMeta<BaseQuery>,
     arg: QueryArg,
   ): ResultType | Promise<ResultType>
@@ -105,7 +106,7 @@ type EndpointDefinitionWithQuery<
   ): unknown
 
   /** A schema for the result *before* it's passed to `transformResponse` */
-  rawResponseSchema?: StandardSchemaV1<BaseQueryResult<BaseQuery>>
+  rawResponseSchema?: StandardSchemaV1<RawResultType>
 
   /** A schema for the error object returned by the `query` or `queryFn`, *before* it's passed to `transformErrorResponse` */
   rawErrorResponseSchema?: StandardSchemaV1<BaseQueryError<BaseQuery>>
@@ -183,10 +184,16 @@ export type BaseEndpointDefinition<
   QueryArg,
   BaseQuery extends BaseQueryFn,
   ResultType,
+  RawResultType extends BaseQueryResult<BaseQuery> = BaseQueryResult<BaseQuery>,
 > = (
   | ([CastAny<BaseQueryResult<BaseQuery>, {}>] extends [NEVER]
       ? never
-      : EndpointDefinitionWithQuery<QueryArg, BaseQuery, ResultType>)
+      : EndpointDefinitionWithQuery<
+          QueryArg,
+          BaseQuery,
+          ResultType,
+          RawResultType
+        >)
   | EndpointDefinitionWithQueryFn<QueryArg, BaseQuery, ResultType>
 ) & {
   /** A schema for the arguments to be passed to the `query` or `queryFn` */
@@ -550,7 +557,8 @@ export type QueryDefinition<
   TagTypes extends string,
   ResultType,
   ReducerPath extends string = string,
-> = BaseEndpointDefinition<QueryArg, BaseQuery, ResultType> &
+  RawResultType extends BaseQueryResult<BaseQuery> = BaseQueryResult<BaseQuery>,
+> = BaseEndpointDefinition<QueryArg, BaseQuery, ResultType, RawResultType> &
   QueryExtraOptions<TagTypes, ResultType, QueryArg, BaseQuery, ReducerPath>
 
 export type InfiniteQueryTypes<
@@ -743,12 +751,14 @@ export type InfiniteQueryDefinition<
   TagTypes extends string,
   ResultType,
   ReducerPath extends string = string,
+  RawResultType extends BaseQueryResult<BaseQuery> = BaseQueryResult<BaseQuery>,
 > =
   // Infinite query endpoints receive `{queryArg, pageParam}`
   BaseEndpointDefinition<
     InfiniteQueryCombinedArg<QueryArg, PageParam>,
     BaseQuery,
-    ResultType
+    ResultType,
+    RawResultType
   > &
     InfiniteQueryExtraOptions<
       TagTypes,
@@ -876,7 +886,8 @@ export type MutationDefinition<
   TagTypes extends string,
   ResultType,
   ReducerPath extends string = string,
-> = BaseEndpointDefinition<QueryArg, BaseQuery, ResultType> &
+  RawResultType extends BaseQueryResult<BaseQuery> = BaseQueryResult<BaseQuery>,
+> = BaseEndpointDefinition<QueryArg, BaseQuery, ResultType, RawResultType> &
   MutationExtraOptions<TagTypes, ResultType, QueryArg, BaseQuery, ReducerPath>
 
 export type EndpointDefinition<
@@ -960,9 +971,21 @@ export type EndpointBuilder<
    *});
    *```
    */
-  query<ResultType, QueryArg>(
+  query<
+    ResultType,
+    QueryArg,
+    RawResultType extends
+      BaseQueryResult<BaseQuery> = BaseQueryResult<BaseQuery>,
+  >(
     definition: OmitFromUnion<
-      QueryDefinition<QueryArg, BaseQuery, TagTypes, ResultType, ReducerPath>,
+      QueryDefinition<
+        QueryArg,
+        BaseQuery,
+        TagTypes,
+        ResultType,
+        ReducerPath,
+        RawResultType
+      >,
       'type'
     >,
   ): QueryDefinition<QueryArg, BaseQuery, TagTypes, ResultType, ReducerPath>
@@ -992,14 +1015,20 @@ export type EndpointBuilder<
    * });
    * ```
    */
-  mutation<ResultType, QueryArg>(
+  mutation<
+    ResultType,
+    QueryArg,
+    RawResultType extends
+      BaseQueryResult<BaseQuery> = BaseQueryResult<BaseQuery>,
+  >(
     definition: OmitFromUnion<
       MutationDefinition<
         QueryArg,
         BaseQuery,
         TagTypes,
         ResultType,
-        ReducerPath
+        ReducerPath,
+        RawResultType
       >,
       'type'
     >,
@@ -1066,27 +1095,31 @@ export function expandTagDescription(
   return typeof description === 'string' ? { type: description } : description
 }
 
-export type QueryArgFrom<D extends BaseEndpointDefinition<any, any, any>> =
-  D extends BaseEndpointDefinition<infer QA, any, any> ? QA : never
+export type QueryArgFrom<D extends BaseEndpointDefinition<any, any, any, any>> =
+  D extends BaseEndpointDefinition<infer QA, any, any, any> ? QA : never
 
 // Just extracting `QueryArg` from `BaseEndpointDefinition`
 // doesn't sufficiently match here.
 // We need to explicitly match against `InfiniteQueryDefinition`
 export type InfiniteQueryArgFrom<
-  D extends BaseEndpointDefinition<any, any, any>,
-> = D extends InfiniteQueryDefinition<infer QA, any, any, any, any> ? QA : never
+  D extends BaseEndpointDefinition<any, any, any, any>,
+> =
+  D extends InfiniteQueryDefinition<infer QA, any, any, any, any, any>
+    ? QA
+    : never
 
 export type QueryArgFromAnyQuery<
-  D extends BaseEndpointDefinition<any, any, any>,
+  D extends BaseEndpointDefinition<any, any, any, any>,
 > =
-  D extends InfiniteQueryDefinition<any, any, any, any, any>
+  D extends InfiniteQueryDefinition<any, any, any, any, any, any>
     ? InfiniteQueryArgFrom<D>
     : D extends QueryDefinition<any, any, any, any>
       ? QueryArgFrom<D>
       : never
 
-export type ResultTypeFrom<D extends BaseEndpointDefinition<any, any, any>> =
-  D extends BaseEndpointDefinition<any, any, infer RT> ? RT : unknown
+export type ResultTypeFrom<
+  D extends BaseEndpointDefinition<any, any, any, any>,
+> = D extends BaseEndpointDefinition<any, any, infer RT, any> ? RT : unknown
 
 export type ReducerPathFrom<
   D extends EndpointDefinition<any, any, any, any, any>,
@@ -1096,9 +1129,11 @@ export type TagTypesFrom<D extends EndpointDefinition<any, any, any, any>> =
   D extends EndpointDefinition<any, any, infer RP, any> ? RP : unknown
 
 export type PageParamFrom<
-  D extends InfiniteQueryDefinition<any, any, any, any, any>,
+  D extends InfiniteQueryDefinition<any, any, any, any, any, any>,
 > =
-  D extends InfiniteQueryDefinition<any, infer PP, any, any, any> ? PP : unknown
+  D extends InfiniteQueryDefinition<any, infer PP, any, any, any, any>
+    ? PP
+    : unknown
 
 export type InfiniteQueryCombinedArg<QueryArg, PageParam> = {
   queryArg: QueryArg

packages/toolkit/src/query/tests/createApi.test-d.ts

@@ -1,6 +1,6 @@
 import { setupApiStore } from '@internal/tests/utils/helpers'
-import type { SerializedError } from '@reduxjs/toolkit'
-import { configureStore } from '@reduxjs/toolkit'
+import type { EntityState, SerializedError } from '@reduxjs/toolkit'
+import { configureStore, createEntityAdapter } from '@reduxjs/toolkit'
 import type {
   DefinitionsFromApi,
   FetchBaseQueryError,
@@ -467,6 +467,7 @@ describe('type tests', () => {
         })
       })
       test('schemas as a source of inference', () => {
+        const postAdapter = createEntityAdapter<Post>()
         const api = createApi({
           baseQuery: fetchBaseQuery({ baseUrl: 'https://example.com' }),
           endpoints: (build) => ({
@@ -482,6 +483,14 @@ describe('type tests', () => {
               argSchema,
               responseSchema: postSchema,
             }),
+            query3: build.query({
+              query: (_arg: void) => `/posts`,
+              rawResponseSchema: v.array(postSchema),
+              transformResponse: (posts) => {
+                expectTypeOf(posts).toEqualTypeOf<Post[]>()
+                return postAdapter.getInitialState(undefined, posts)
+              },
+            }),
           }),
         })
 
@@ -496,6 +505,11 @@ describe('type tests', () => {
         expectTypeOf(
           api.endpoints.query2.Types.ResultType,
         ).toEqualTypeOf<Post>()
+
+        expectTypeOf(api.endpoints.query3.Types.QueryArg).toEqualTypeOf<void>()
+        expectTypeOf(api.endpoints.query3.Types.ResultType).toEqualTypeOf<
+          EntityState<Post, Post['id']>
+        >()
       })
     })
   })