Leandro Facchinetti
Leandro Facchinetti

Reputation: 108

How do I type a function that works over a Discriminated Union and knows that the type of the input is the same as the type of the output?

Using TypeScript, how do I type a function that works over a Discriminated Union and knows that the type of the input is the same as the type of the output?

Example

Suppose that I have types like those in the Handbook:

type Shape = Square | Rectangle | Circle;
interface Square {
  kind: "square";
  size: number;
}
interface Rectangle {
  kind: "rectangle";
  width: number;
  height: number;
}
interface Circle {
  kind: "circle";
  radius: number;
}

I want to define a function that works over any kind of Shape and preserves that kind (it won’t, say, receive a Square and return a Circle), for example:

function double(s) {
  switch (s.kind) {
    case "square":
      return { ...s, size: s.size * 2 };
    case "rectangle":
      return { ...s, height: s.height * 2, width: s.width * 2 };
    case "circle":
      return { ...s, radius: s.radius * 2 };
  }
}

In particular, I want something like the following to work:

const square: Square = { kind: "square", size: 4 };
const doubleSquare: Square = double(square);

Here are some things that I tried and didn’t work:

Let the Type Be Shape

function double(s: Shape): Shape {
  // ...
}

The problem with this approach is that it loses the connection between input and output: as far as the type system is concerned, double could received a Square and return a Circle. Then the type system rightfully complains that doubleSquare may be any kind of Shape.

Use Generic Constraints

function double<T extends Shape>(s: T): T {
  // ...
}

Now the type system is happy about doubleSquare, but it doesn’t like the Type Guards in the switch, for example in

case "square":
  return { ...s, size: s.size * 2 };

it complains about s.size saying that s still has type T extends Shape.

I expected this to work, particularly because Shape is an Union Type—there are only so many kinds it may be (Square, Rectangle, or Circle).

Overload the Function

function double(s: Square): Square;
function double(s: Rectangle): Rectangle;
function double(s: Circle): Circle;
function double(s: Shape): Shape {
  // ...
}

This actually makes the type system happy, but it makes me sad. First, because I have to spell out all kinds of Shape. Second, because the type system won’t help me if I make a mistake in the implementation, for example, if I receive a Square and return a Circle:

case "square":
  return { kind: "circle", radius: s.size * 2 };

The Workaround

Currently, I’m using the type described in Let the Type Be Shape and using a Type Assertion to silence the type error:

const square: Square = { kind: "square", size: 4 };
const doubleSquare: Square = double(square) as typeof square;

Other Things I Heard About

After some research, I think that maybe what I need is a Generalized Algebraic Data Type (GADT), or a Conditional Type. But these seem to be big hammers, and I don’t even know exactly how to use them.

Upvotes: 2

Views: 77

Answers (1)

Titian Cernicova-Dragomir
Titian Cernicova-Dragomir

Reputation: 249666

A combination of the generic and overloaded approaches work best here. The generic signature can be the public signature, this will ensure that the type passed in is the same as the type passed out. The function double(s: Shape): Shape will be the implementation signature and will allow us to manipulate the parameter easier:

function double<T extends Shape>(s: T): T
function double(s: Shape) : Shape {
  switch (s.kind) {
    case "square":
      return { ...s, size: s.size * 2 };
    case "rectangle":
      return { ...s, height: s.height * 2, width: s.width * 2 };
    case "circle":
      return { ...s, radius: s.radius * 2 };
  }
}

Playground

This will not ensure the implementation itself returns the same type as passed in, the developer it still responsible for that. I have a workaround in a similar question

Upvotes: 1

Related Questions