add types for manually typing hook results in userland code (#2276)

Author: phryneas

packages/toolkit/src/query/react/buildHooks.ts

@@ -52,6 +52,7 @@ import { useStableQueryArgs } from './useSerializedStableValue'
 import type { UninitializedValue } from './constants'
 import { UNINITIALIZED_VALUE } from './constants'
 import { useShallowStableValue } from './useShallowStableValue'
+import type { BaseQueryFn } from '../baseQueryTypes'
 
 // Copy-pasted from React-Redux
 export const useIsomorphicLayoutEffect =
@@ -97,7 +98,26 @@ export type UseQuery<D extends QueryDefinition<any, any, any, any>> = <
 >(
   arg: QueryArgFrom<D> | SkipToken,
   options?: UseQuerySubscriptionOptions & UseQueryStateOptions<D, R>
-) => UseQueryStateResult<D, R> & ReturnType<UseQuerySubscription<D>>
+) => UseQueryHookResult<D, R>
+
+export type UseQueryHookResult<
+  D extends QueryDefinition<any, any, any, any>,
+  R = UseQueryStateDefaultResult<D>
+> = UseQueryStateResult<D, R> & UseQuerySubscriptionResult<D>
+
+/**
+ * Helper type to manually type the result
+ * of the `useQuery` hook in userland code.
+ */
+export type TypedUseQueryHookResult<
+  ResultType,
+  QueryArg,
+  BaseQuery extends BaseQueryFn,
+  R = UseQueryStateDefaultResult<
+    QueryDefinition<QueryArg, BaseQuery, string, ResultType, string>
+  >
+> = TypedUseQueryStateResult<ResultType, QueryArg, BaseQuery, R> &
+  TypedUseQuerySubscriptionResult<ResultType, QueryArg, BaseQuery>
 
 interface UseQuerySubscriptionOptions extends SubscriptionOptions {
   /**
@@ -162,7 +182,23 @@ export type UseQuerySubscription<
 > = (
   arg: QueryArgFrom<D> | SkipToken,
   options?: UseQuerySubscriptionOptions
-) => Pick<QueryActionCreatorResult<D>, 'refetch'>
+) => UseQuerySubscriptionResult<D>
+
+export type UseQuerySubscriptionResult<
+  D extends QueryDefinition<any, any, any, any>
+> = Pick<QueryActionCreatorResult<D>, 'refetch'>
+
+/**
+ * Helper type to manually type the result
+ * of the `useQuerySubscription` hook in userland code.
+ */
+export type TypedUseQuerySubscriptionResult<
+  ResultType,
+  QueryArg,
+  BaseQuery extends BaseQueryFn
+> = UseQuerySubscriptionResult<
+  QueryDefinition<QueryArg, BaseQuery, string, ResultType, string>
+>
 
 export type UseLazyQueryLastPromiseInfo<
   D extends QueryDefinition<any, any, any, any>
@@ -338,6 +374,19 @@ export type UseQueryStateResult<
   R
 > = NoInfer<R>
 
+/**
+ * Helper type to manually type the result
+ * of the `useQueryState` hook in userland code.
+ */
+export type TypedUseQueryStateResult<
+  ResultType,
+  QueryArg,
+  BaseQuery extends BaseQueryFn,
+  R = UseQueryStateDefaultResult<
+    QueryDefinition<QueryArg, BaseQuery, string, ResultType, string>
+  >
+> = NoInfer<R>
+
 type UseQueryStateBaseResult<D extends QueryDefinition<any, any, any, any>> =
   QuerySubState<D> & {
     /**
@@ -435,6 +484,22 @@ export type UseMutationStateResult<
   reset: () => void
 }
 
+/**
+ * Helper type to manually type the result
+ * of the `useMutation` hook in userland code.
+ */
+export type TypedUseMutationResult<
+  ResultType,
+  QueryArg,
+  BaseQuery extends BaseQueryFn,
+  R = MutationResultSelectorResult<
+    MutationDefinition<QueryArg, BaseQuery, string, ResultType, string>
+  >
+> = UseMutationStateResult<
+  MutationDefinition<QueryArg, BaseQuery, string, ResultType, string>,
+  R
+>
+
 /**
  * A React hook that lets you trigger an update request for a given endpoint, and subscribes the component to read the request status from the Redux store. The component will re-render as the loading status changes.
  *

packages/toolkit/src/query/react/index.ts

@@ -21,4 +21,10 @@ const createApi = /* @__PURE__ */ buildCreateApi(
   reactHooksModule()
 )
 
+export type {
+  TypedUseQueryHookResult,
+  TypedUseQueryStateResult,
+  TypedUseQuerySubscriptionResult,
+  TypedUseMutationResult,
+} from './buildHooks'
 export { createApi, reactHooksModule }

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

@@ -1,10 +1,17 @@
 import type { SerializedError } from '@reduxjs/toolkit'
-import type { FetchBaseQueryError } from '@reduxjs/toolkit/query/react'
+import type {
+  FetchBaseQueryError,
+  TypedUseQueryHookResult,
+  TypedUseQueryStateResult,
+  TypedUseQuerySubscriptionResult,
+  TypedUseMutationResult,
+} from '@reduxjs/toolkit/query/react'
 import { createApi, fetchBaseQuery } from '@reduxjs/toolkit/query/react'
 import { expectExactType, expectType } from './helpers'
 
+const baseQuery = fetchBaseQuery()
 const api = createApi({
-  baseQuery: fetchBaseQuery(),
+  baseQuery,
   endpoints: (build) => ({
     test: build.query<string, void>({ query: () => '' }),
     mutation: build.mutation<string, void>({ query: () => '' }),
@@ -547,4 +554,52 @@ describe.skip('TS only tests', () => {
       expectType<never>(result)
     }
   })
+
+  test('"Typed" helper types', () => {
+    // useQuery
+    {
+      const result = api.endpoints.test.useQuery()
+      expectType<TypedUseQueryHookResult<string, void, typeof baseQuery>>(
+        result
+      )
+    }
+    // useQuery with selectFromResult
+    {
+      const result = api.endpoints.test.useQuery(undefined, {
+        selectFromResult: () => ({ x: true }),
+      })
+      expectType<
+        TypedUseQueryHookResult<string, void, typeof baseQuery, { x: boolean }>
+      >(result)
+    }
+    // useQueryState
+    {
+      const result = api.endpoints.test.useQueryState()
+      expectType<TypedUseQueryStateResult<string, void, typeof baseQuery>>(
+        result
+      )
+    }
+    // useQueryState with selectFromResult
+    {
+      const result = api.endpoints.test.useQueryState(undefined, {
+        selectFromResult: () => ({ x: true }),
+      })
+      expectType<
+        TypedUseQueryStateResult<string, void, typeof baseQuery, { x: boolean }>
+      >(result)
+    }
+    // useQuerySubscription
+    {
+      const result = api.endpoints.test.useQuerySubscription()
+      expectType<
+        TypedUseQuerySubscriptionResult<string, void, typeof baseQuery>
+      >(result)
+    }
+
+    // useMutation
+    {
+      const [trigger, result] = api.endpoints.mutation.useMutation()
+      expectType<TypedUseMutationResult<string, void, typeof baseQuery>>(result)
+    }
+  })
 })