Conditional return from a function based on conditional parameter

Question

A function that has a parameter with a conditional type should return conditionally based on that. Is that possible?

Here's a playground

---

I have a function that takes one parameter, which is either a custom type (QueryKey) or a function:

export function createHandler(
  createQueryKeyOrQueryKey: QueryKey | ((idForQueryKey: string) => QueryKey),
) { ... }

Based on this parameter, I need to return different types from the createHandler function:

  return {
    createState:
      typeof createQueryKeyOrQueryKey !== "function"
        ? (data) => createState(data, createQueryKeyOrQueryKey)
        : (data, id) => createState(data, createQueryKeyOrQueryKey(id)),
  };

What is returned with createState is also conditional:

createState:
  | ((data: TData) => State)
  | ((data: TData, idForQueryKey: string) => State);

When creating a handler via createHandler, there should be two ways to use it:

handler.createState({ a: 123 })
handler.createState({ a: 123 }, "some-id")

...but only one variant of createState should be allowed at a time. Which one is allowed should be based on how the handler is created, where either a query key OR a function is used:

// Option 1:
// 
// query key directly
const queryKey = "some-query-key"
const handler1 = createHandler(queryKey)
// ✅ should be allowed
handler1.createState({ a: 123 })
// ❌ should not be allowed
handler1.createState({ a: 123 }, "some-id")

// Option 2:
//
// query key as function for creation
const queryKeyCreator = (id: string) => "some-query-key" + id
const handler2 = createHandler(queryKeyCreator)
// ❌ should not be allowed
handler2.createState({ a: 123 })
// ✅ should be allowed
handler2.createState({ a: 123 }, "some-id")

Right now, the return type does not work correctly (data is any):

Why so? TypeScript knows that createState is conditional and the one where the return function only has one parameter (data) is also a valid type based on the type of createState.

---

BTW: Is there maybe a better way to solve this problem? Function overloading or discriminating unions via key would maybe work, but I'm not 100% sure how the different implementations can be based on the caller's decision which variant (key or function) is used.

Answer 1

What is returned with createState is also conditional

It's not. It's a union, it basically says "createState can either return this kind of function, or this one, if you call createState then you will have to figure it out by yourself". What you write in implementation doesn't matter. Basically if you provide the return type yourself (that is, typescript doesn't have to infer it based on the implementation of the function), typescript just sees this:

function createHandler<TData extends QueryData>(
  createQueryKeyOrQueryKey: ((idForQueryKey: string) => QueryKey) | QueryKey,
): {
  createState: 
   | ((data: TData, idForQueryKey: string) => State)
   | ((data: TData) => State);
} {
  // Stuff
}

It doesn't care about implementation and has no idea that you want createQueryKeyOrQueryKey and createState to be related in any way. Conditional types are actually types that look like ... extends ... ? ... : ....

Your case is the exact case that overloads were designed for. It should look something like this

// First possible signature
// If the caller passes in a function, we return variant with `idForQuery`
function createHandler<TData extends QueryData>(
  createQueryKey: (idForQueryKey: string) => QueryKey
): {
  createState: (data: TData, idForQueryKey: string) => State
}
// Second signature
// If the caller passes in a string, we return the simpler variant
function createHandler<TData extends QueryData>(
  queryKey: QueryKey
): {
  createState: (data: TData) => State
}
// Implementation signature
// This is not visible to the callers of your function, it only exists
// so that you can type everything correctly inside the function body
function createHandler<TData extends QueryData>(
  createQueryKeyOrQueryKey: ((idForQueryKey: string) => QueryKey) | QueryKey,
): {
  createState: 
   | ((data: TData, idForQueryKey: string) => State)
   | ((data: TData) => State);
} {
  return {
    createState:
      typeof createQueryKeyOrQueryKey !== "function"
        ? (data: TData) => createState(data, createQueryKeyOrQueryKey)
        : (data, id) =>
          createState(data, createQueryKeyOrQueryKey(id)),
  };
}

See sandbox

---

There is also a way to do it using actual conditional types, but for this you will need the second generic type argument to keep track of what type is passed to createQueryKeyOrQueryKey:

declare function createHandler<
  TData extends QueryData,
  TQueryKey extends ((idForQueryKey: string) => QueryKey) | QueryKey
>(
  createQueryKeyOrQueryKey: TQueryKey
): {
  createState:
    // Depending on the type of the argument, choose the appropriate variant
    TQueryKey extends QueryKey ?
      (data: TData) => State :
      (data: TData, idForQueryKey: string) => State
}

See sandbox

But you most definitely should not use this, because

1. I wrote declare function and didn't write the implementation for a good reason: it will probably be a headache to type properly, you will most likely have to use type casts
2. There are all sorts of weird edge cases, like you can do createHandler<..., any>(...) for example and you will lose a good part of your types
3. You now have 2 generics, which means it will be hard to use, because in TS you either don't specify any generics, or you specify all of them, there is no in between

This way of doing it is just so you know, but overloads are the perfect solution so you should definitely use those