Conditionally prohibit function parameter depending on enum value

I have a function fn with two possible call signatures:

enum Mode {
  NEEDS_ARG,
  ANYTHING_ELSE
}

fn(Mode.NEEDS_ARG, arg)  // Call signature 1
fn(Mode.ANYTHING_ELSE)   // Call signature 2

fn(Mode.ANYTHING_ELSE, arg)   // Invalid

Depending on the value of the first argument, I'd like to prohibit the second argument.

Here's my attempt at typing such a function:

enum Mode {
  NEEDS_ARG,
  ANYTHING_ELSE
}

type Arg = "argument"

function fn(
  mode: Mode,
  arg: typeof mode extends Mode.NEEDS_ARG ? Arg : never
) {
  arg  // TS infers this as 'never'
}

From the Typescript docs on Conditional Types:

When the type on the left of the extends is assignable to the one on the right, then you’ll get the type in the first branch (the “true” branch); otherwise you’ll get the type in the latter branch (the “false” branch).

I see why my implementation doesn't work - typeof mode evaluates to Mode, which is not assignable to Mode.NEEDS_ARG.

But what is the right way to type this function?

EDIT

I run into a couple of issues trying to adapt the generic argument approach provided by @jcalz

Issue 1: arg type doesn't get narrowed based on the value of mode

Playground

I expected the type of arg to be narrowed when mode === Mode.NEEDS_ARG, but it isn't

function fn<T extends Mode>(
  mode: T,
  ...rest: T extends Mode.NEEDS_ARG ? [arg: Arg] : []
) {
  const [arg] = rest
  if(arg !== undefined) {
    arg  // Type is narrowed to 'Arg'
  }
  if(mode === Mode.NEEDS_ARG) {
    arg  // Typescript still thinks this is 'Arg | undefined'
  }
}

Issue 2: Doesn't work without the spread syntax

Playground

If I replace the ...rest param with a single arg param, Typescript complains

function fn<T extends Mode>(
  mode: T,
  arg: T extends Mode.NEEDS_ARG ? Arg: never
) {}

fn(Mode.NEEDS_ARG, arg);  // Ok
fn(Mode.NO_ARG);          // Err: 'An argument for 'arg' was not provided.'

Solution 1:

So you want to allow these calls:

fn(Mode.NEEDS_ARG, arg); // okay
fn(Mode.ANYTHING_ELSE); // okay

While disallowing these:

fn(Mode.NEEDS_ARG); // error!
fn(Mode.ANYTHING_ELSE, arg); // error!

There are several ways to accomplish this in TypeScript.

The traditional solution is to make fn() an overloaded function with two call signatures:

// call signatures
function fn(mode: Mode.NEEDS_ARG, arg: Arg): void;
function fn(mode: Mode.ANYTHING_ELSE): void;
// impl
function fn(mode: Mode, arg?: Arg) {

}

Nowadays if you have overloads that have the same return type, you can refactor overloads into a rest parameter whose type is a union of tuples:

function fn(
  ...args: [mode: Mode.NEEDS_ARG, arg: Arg] | [mode: Mode.ANYTHING_ELSE]
): void {

}

Alternatively you could make the function generic and use a conditional type to allow/disallow a second argument based on the specific type of the first argument:

function fn<T extends Mode>(
  mode: T,
  ...rest: T extends Mode.NEEDS_ARG ? [arg: Arg] : []
): void {

}

Because the number of function parameters changes depending on T, you can't rewrite this as a single call signature without a rest parameter. There was some work done to let you mark a parameter as type void and it would be optional, but it doesn't work with generics (see microsoft/TypeScript#29131), so it's just not possible. You need a rest parameter to write the call signature this way.

All of those should work to some degree from the caller's side. This is more or less the end of the answer to the question as originally asked.

In the implementation, though, there are limitations. If you really want to treat the arguments as a discriminated union and narrow based on whether mode is Mode.NEEDS_ARG or not, and you want type safety from the compiler, you pretty much need to use the union-of-tuples technique:

function fn(
  ...args: [mode: Mode.NEEDS_ARG, arg: Arg] | [mode: Mode.ANYTHING_ELSE]
): void {
  const mode = args[0];
  if (mode === Mode.NEEDS_ARG) {
    const arg = args[1];
    arg.toUpperCase(); // okay
  } else {
    args.length // known to be 1
  }
}

Inside the function we get a variable named mode from args, such that when we check mode it acts as a discriminant for args (this support was added in TS4.3; before this you would need to say args[0] everywhere instead of mode).

It's not beautiful but it works.

If you do try this with the overloads, your implementation will be too loosely typed and it will fail to realize that checking mode has any implication for the type of arg

// call signatures
function fn(mode: Mode.NEEDS_ARG, arg: Arg): void;
function fn(mode: Mode.ANYTHING_ELSE): void;
// impl
function fn(mode: Mode, arg?: Arg) {
  if (mode === Mode.NEEDS_ARG) {
    arg.toUpperCase(); // error
  }
}

You can of course use a type assertion like arg!.toUpperCase(), but this is not very type safe. Overload implementations are sort of "cut off" from the call signatures, and do not "prune" the call signature set based on control flow. See microsoft/TypeScript#22609 for a relevant feature request to support this.

On the other hand, if you try to do this with the generic conditional, you have a similar problem; in this case it's because the compiler does not know what T is inside the body, and thus does not know how to evaluate T extends Mode.NEEDS_ARG ? [arg: Arg] : []:

function fn<T extends Mode>(
  mode: T,
  ...rest: T extends Mode.NEEDS_ARG ? [arg: Arg] : []
): void {
  rest // T extends Mode.NEEDS_ARG ? [arg: "argument"] : []
  if (mode === Mode.NEEDS_ARG) {
    rest[1].toUpperCase(); // error
  }
}

Again, you can use type assertions or other ways of suppressing the error, but currently the body of generic functions does not know how to do much with conditional types. You just can't narrow the type of a generic type parameter by checking its value. See https://github.com/microsoft/TypeScript/issues/27808 for a feature request that could help with this.

So there you go. You have several options for giving types to the call signature(s) for the function. If you need to have type safety inside the implementation, the closest you can get is to use the union-of-rest-tuple approach. Otherwise you will have to forgo some type safety.

Playground link to code