enable enhanceEndpoints.transformResponse to override ResultType

packages/toolkit/src/query/apiTypes.ts

@@ -2,7 +2,7 @@ import type {
   EndpointDefinitions,
   EndpointBuilder,
   EndpointDefinition,
-  ReplaceTagTypes,
+  UpdateDefinitions,
 } from './endpointDefinitions'
 import type {
   UnionToIntersection,
@@ -93,11 +93,15 @@ export type Api<
   /**
    *A function to enhance a generated API with additional information. Useful with code-generation.
    */
-  enhanceEndpoints<NewTagTypes extends string = never>(_: {
+  enhanceEndpoints<
+    NewTagTypes extends string = never,
+    NewDefinitions extends EndpointDefinitions = never
+  >(_: {
     addTagTypes?: readonly NewTagTypes[]
-    endpoints?: ReplaceTagTypes<
+    endpoints?: UpdateDefinitions<
       Definitions,
-      TagTypes | NoInfer<NewTagTypes>
+      TagTypes | NoInfer<NewTagTypes>,
+      NewDefinitions
     > extends infer NewDefinitions
       ? {
           [K in keyof NewDefinitions]?:
@@ -107,7 +111,7 @@ export type Api<
       : never
   }): Api<
     BaseQuery,
-    ReplaceTagTypes<Definitions, TagTypes | NewTagTypes>,
+    UpdateDefinitions<Definitions, TagTypes | NewTagTypes, NewDefinitions>,
     ReducerPath,
     TagTypes | NewTagTypes,
     Enhancers

packages/toolkit/src/query/endpointDefinitions.ts

@@ -16,8 +16,10 @@ import type {
   MaybePromise,
   OmitFromUnion,
   CastAny,
+  NonUndefined,
 } from './tsHelpers'
 import type { NEVER } from './fakeBaseQuery'
+import type { Api } from '@reduxjs/toolkit/query'
 
 const resultType = /* @__PURE__ */ Symbol()
 const baseQuery = /* @__PURE__ */ Symbol()
@@ -775,9 +777,58 @@ export type ReducerPathFrom<
 export type TagTypesFrom<D extends EndpointDefinition<any, any, any, any>> =
   D extends EndpointDefinition<any, any, infer RP, any> ? RP : unknown
 
-export type ReplaceTagTypes<
+export type TagTypesFromApi<T> = T extends Api<any, any, any, infer TagTypes>
+  ? TagTypes
+  : never
+
+export type DefinitionsFromApi<T> = T extends Api<
+  any,
+  infer Definitions,
+  any,
+  any
+>
+  ? Definitions
+  : never
+
+export type TransformedResponse<
+  NewDefinitions extends EndpointDefinitions,
+  K,
+  ResultType
+> = K extends keyof NewDefinitions
+  ? NewDefinitions[K]['transformResponse'] extends undefined
+    ? ResultType
+    : ReturnType<NonUndefined<NewDefinitions[K]['transformResponse']>>
+  : ResultType
+
+export type OverrideResultType<Definition, NewResultType> =
+  Definition extends QueryDefinition<
+    infer QueryArg,
+    infer BaseQuery,
+    infer TagTypes,
+    any,
+    infer ReducerPath
+  >
+    ? QueryDefinition<QueryArg, BaseQuery, TagTypes, NewResultType, ReducerPath>
+    : Definition extends MutationDefinition<
+        infer QueryArg,
+        infer BaseQuery,
+        infer TagTypes,
+        any,
+        infer ReducerPath
+      >
+    ? MutationDefinition<
+        QueryArg,
+        BaseQuery,
+        TagTypes,
+        NewResultType,
+        ReducerPath
+      >
+    : never
+
+export type UpdateDefinitions<
   Definitions extends EndpointDefinitions,
-  NewTagTypes extends string
+  NewTagTypes extends string,
+  NewDefinitions extends EndpointDefinitions
 > = {
   [K in keyof Definitions]: Definitions[K] extends QueryDefinition<
     infer QueryArg,
@@ -786,7 +837,13 @@ export type ReplaceTagTypes<
     infer ResultType,
     infer ReducerPath
   >
-    ? QueryDefinition<QueryArg, BaseQuery, NewTagTypes, ResultType, ReducerPath>
+    ? QueryDefinition<
+        QueryArg,
+        BaseQuery,
+        NewTagTypes,
+        TransformedResponse<NewDefinitions, K, ResultType>,
+        ReducerPath
+      >
     : Definitions[K] extends MutationDefinition<
         infer QueryArg,
         infer BaseQuery,
@@ -798,7 +855,7 @@ export type ReplaceTagTypes<
         QueryArg,
         BaseQuery,
         NewTagTypes,
-        ResultType,
+        TransformedResponse<NewDefinitions, K, ResultType>,
         ReducerPath
       >
     : never

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

@@ -1,11 +1,15 @@
 import { configureStore, createAction, createReducer } from '@reduxjs/toolkit'
+import type { SerializedError } from '@reduxjs/toolkit'
 import type {
   Api,
   MutationDefinition,
   QueryDefinition,
 } from '@reduxjs/toolkit/query'
 import { createApi, fetchBaseQuery } from '@reduxjs/toolkit/query'
-import type { FetchBaseQueryMeta } from '@reduxjs/toolkit/dist/query/fetchBaseQuery'
+import type {
+  FetchBaseQueryError,
+  FetchBaseQueryMeta,
+} from '@reduxjs/toolkit/dist/query/fetchBaseQuery'
 
 import {
   ANY,
@@ -19,6 +23,11 @@ import { server } from './mocks/server'
 import { rest } from 'msw'
 import type { SerializeQueryArgs } from '../defaultSerializeQueryArgs'
 import { string } from 'yargs'
+import type {
+  DefinitionsFromApi,
+  OverrideResultType,
+  TagTypesFromApi,
+} from '@reduxjs/toolkit/dist/query/endpointDefinitions'
 
 const originalEnv = process.env.NODE_ENV
 beforeAll(() => void ((process.env as any).NODE_ENV = 'development'))
@@ -522,6 +531,71 @@ describe('endpoint definition typings', () => {
         ['modified2', { ...commonBaseQueryApi, forced: undefined }, undefined],
       ])
     })
+
+    test('updated transform response types', async () => {
+      const baseApi = createApi({
+        baseQuery: fetchBaseQuery({ baseUrl: 'https://example.com' }),
+        tagTypes: ['old'],
+        endpoints: (build) => ({
+          query1: build.query<'out1', void>({ query: () => 'success' }),
+          mutation1: build.mutation<'out1', void>({ query: () => 'success' }),
+        }),
+      })
+
+      type Transformed = { value: string }
+
+      type Definitions = DefinitionsFromApi<typeof api>
+      type TagTypes = TagTypesFromApi<typeof api>
+
+      type Q1Definition = OverrideResultType<Definitions['query1'], Transformed>
+      type M1Definition = OverrideResultType<
+        Definitions['mutation1'],
+        Transformed
+      >
+
+      type UpdatedDefitions = Omit<Definitions, 'query1' | 'mutation1'> & {
+        query1: Q1Definition
+        mutation1: M1Definition
+      }
+
+      const enhancedApi = baseApi.enhanceEndpoints<TagTypes, UpdatedDefitions>({
+        endpoints: {
+          query1: {
+            transformResponse: (a, b, c) => ({
+              value: 'transformed',
+            }),
+          },
+          mutation1: {
+            transformResponse: (a, b, c) => ({
+              value: 'transformed',
+            }),
+          },
+        },
+      })
+
+      const storeRef = setupApiStore(enhancedApi, undefined, {
+        withoutTestLifecycles: true,
+      })
+
+      const queryResponse = await storeRef.store.dispatch(
+        enhancedApi.endpoints.query1.initiate()
+      )
+      expect(queryResponse.data).toEqual({ value: 'transformed' })
+      expectType<Transformed | Promise<Transformed> | undefined>(
+        queryResponse.data
+      )
+
+      const mutationResponse = await storeRef.store.dispatch(
+        enhancedApi.endpoints.mutation1.initiate()
+      )
+      expectType<
+        | { data: Transformed | Promise<Transformed> }
+        | { error: FetchBaseQueryError | SerializedError }
+      >(mutationResponse)
+      expect('data' in mutationResponse && mutationResponse.data).toEqual({
+        value: 'transformed',
+      })
+    })
   })
 })
 

packages/toolkit/src/query/tsHelpers.ts

@@ -32,6 +32,8 @@ export type OptionalIfAllPropsOptional<T> = HasRequiredProps<T, T, T | never>
 
 export type NoInfer<T> = [T][T extends any ? 0 : never]
 
+export type NonUndefined<T> = T extends undefined ? never : T
+
 export type UnwrapPromise<T> = T extends PromiseLike<infer V> ? V : T
 
 export type MaybePromise<T> = T | PromiseLike<T>