types: improve useSWRSubscription types

subscription/index.ts

@@ -8,19 +8,34 @@ import {
   createCacheHelper
 } from 'swr/_internal'
 
-export type SWRSubscription<Data = any, Error = any> = (
-  key: Key,
-  { next }: { next: (err?: Error | null, data?: Data) => void }
-) => () => void
+export type SWRSubNext<Data = any, Error = any> = {
+  next: (err?: Error | null, data?: Data) => void
+}
+
+export type SWRSubscription<
+  SWRSubKey extends Key = Key,
+  Data = any,
+  Error = any
+> = SWRSubKey extends () => infer Arg | null | undefined | false
+  ? (key: Arg, { next }: SWRSubNext<Data, Error>) => void
+  : SWRSubKey extends null | undefined | false
+  ? never
+  : SWRSubKey extends infer Arg
+  ? (key: Arg, { next }: SWRSubNext<Data, Error>) => void
+  : never
 
 export type SWRSubscriptionResponse<Data = any, Error = any> = {
   data?: Data
   error?: Error
 }
 
-export type SWRSubscriptionHook<Data = any, Error = any> = (
-  key: Key,
-  subscribe: SWRSubscription<Data, Error>,
+export type SWRSubscriptionHook = <
+  Data = any,
+  Error = any,
+  SWRSubKey extends Key = Key
+>(
+  key: SWRSubKey,
+  subscribe: SWRSubscription<SWRSubKey, Data, Error>,
   config?: SWRConfiguration
 ) => SWRSubscriptionResponse<Data, Error>
 
@@ -30,10 +45,10 @@ const subscriptionStorage = new WeakMap<object, SubscriptionStates>()
 
 const SUBSCRIPTION_PREFIX = '$sub$'
 
-export const subscription = (<Data, Error>(useSWRNext: SWRHook) =>
+export const subscription = (<Data = any, Error = any>(useSWRNext: SWRHook) =>
   (
     _key: Key,
-    subscribe: SWRSubscription<Data, Error>,
+    subscribe: SWRSubscription<any, Data, Error>,
     config: SWRConfiguration & typeof SWRConfig.defaultValue
   ): SWRSubscriptionResponse<Data, Error> => {
     const [key] = serialize(_key)

test/type/sub.ts

@@ -0,0 +1,95 @@
+import useSWRSubscription from 'swr/subscription'
+import type { SWRSubNext, SWRSubscription } from 'swr/subscription'
+import { expectType, truthy } from './utils'
+
+export function useTestSubscription() {
+  useSWRSubscription('key', (key, { next: _ }: SWRSubNext<string, Error>) => {
+    expectType<'key'>(key)
+    return () => {}
+  })
+  useSWRSubscription(
+    truthy() ? 'key' : undefined,
+    (key, { next: _ }: SWRSubNext<string, Error>) => {
+      expectType<'key'>(key)
+      return () => {}
+    }
+  )
+  useSWRSubscription(
+    ['key', 1],
+    (key, { next: _ }: SWRSubNext<string, Error>) => {
+      expectType<[string, number]>(key)
+      return () => {}
+    }
+  )
+  useSWRSubscription(
+    truthy() ? ['key', 1] : undefined,
+    (key, { next: _ }: SWRSubNext<string, Error>) => {
+      expectType<[string, number]>(key)
+      return () => {}
+    }
+  )
+  useSWRSubscription(
+    { foo: 'bar' },
+    (key, { next: _ }: SWRSubNext<string, Error>) => {
+      expectType<{ foo: string }>(key)
+      return () => {}
+    }
+  )
+  useSWRSubscription(
+    truthy() ? { foo: 'bar' } : undefined,
+    (key, { next: _ }: SWRSubNext<string, Error>) => {
+      expectType<{ foo: string }>(key)
+      return () => {}
+    }
+  )
+
+  useSWRSubscription(
+    () => 'key',
+    (key, { next: _ }: SWRSubNext<string, Error>) => {
+      expectType<string>(key)
+      return () => {}
+    }
+  )
+  useSWRSubscription(
+    () => (truthy() ? 'key' : undefined),
+    (key, { next: _ }: SWRSubNext<string, Error>) => {
+      expectType<'key'>(key)
+      return () => {}
+    }
+  )
+  useSWRSubscription(
+    () => ['key', 1],
+    (key, { next: _ }: SWRSubNext<string, Error>) => {
+      expectType<[string, number]>(key)
+      return () => {}
+    }
+  )
+  useSWRSubscription(
+    () => (truthy() ? ['key', 1] : undefined),
+    (key, { next: _ }: SWRSubNext<string, Error>) => {
+      expectType<[string, number]>(key)
+      return () => {}
+    }
+  )
+  useSWRSubscription(
+    () => ({ foo: 'bar' }),
+    (key, { next: _ }: SWRSubNext<string, Error>) => {
+      expectType<{ foo: string }>(key)
+      return () => {}
+    }
+  )
+  useSWRSubscription(
+    () => (truthy() ? { foo: 'bar' } : undefined),
+    (key, { next: _ }: SWRSubNext<string, Error>) => {
+      expectType<{ foo: string }>(key)
+      return () => {}
+    }
+  )
+
+  const sub: SWRSubscription<string, string, Error> = (_, { next: __ }) => {
+    return () => {}
+  }
+  const { data: data2, error: error2 } = useSWRSubscription('key', sub)
+  expectType<string | undefined>(data2)
+  expectType<Error | undefined>(error2)
+}