TypeScript conditional return type: works for one condition but not for two

Question

I have a factory function that creates shapes based on a discriminated union of shape arguments, for example:

interface CircleArgs { type: "circle", radius: number };
interface SquareArgs { type: "square", length: number };

type ShapeArgs = CircleArgs | SquareArgs;

class Circle { constructor(_: CircleArgs) {}}
class Square { constructor(_: SquareArgs) {}}

type Shape = Circle | Square;

type ShapeOfArgs<Args extends ShapeArgs> = 
   Args extends CircleArgs? Circle :
   Square;

function createShape<Args extends ShapeArgs>(args: Args): ShapeOfArgs<Args> {
    switch (args.type) {
        case "circle": return new Circle(args);
        case "square": return new Square(args);
    }
}

Doing it this way not only helps TS to infer the right return type from the argument:

const circle1 = createShape({type: "circle", radius: 1}); // inferred as Circle
const square1 = createShape({type: "square", length: 1}); // inferred as Square

But also correctly propagates the return type if I need to wrap the call in another function:

class Container { insert(_: Shape): void {}; }

function createShapeIn<Args extends ShapeArgs>(args: Args, cont: Container) {
    const shape = createShape(args);
    cont.insert(shape);

    return shape;
}

const cont = new Container();
const circle2 = createShapeIn({type: "circle", radius: 1}, cont); // still inferred as Circle
const square2 = createShapeIn({type: "square", length: 1}, cont); // still inferred as Square

However...

If I add just one more type to the mix everything breaks.

Let's add another type of shapes called Figure:

interface FigureArgs { type: "figure", amount: number };
class Figure { constructor(_: FigureArgs) {}}

Then the our union types will be:

type ShapeArgs = CircleArgs | SquareArgs | FigureArgs;
type Shape = Circle | Square | Figure;

Finally, the conditional type is changed to:

type ShapeOfArgs<Args extends ShapeArgs> = 
    Args extends CircleArgs? Circle :
    Args extends SquareArgs? Square :
    Figure;

And now this no longer transpiles:

function createShape<Args extends ShapeArgs>(args: Args): ShapeOfArgs<Args> {
    switch (args.type) {
        case "circle": return new Circle(args); // error: Circle not assignable ShapeOfArgs<Args>
        case "square": return new Square(args); // error: Square not assignable ShapeOfArgs<Args>
        case "figure": return new Figure(args); // error: Figure not assignable ShapeOfArgs<Args>
    }
}

What am I doing wrong?

PS: I know I can force it using the as operator, but that defeats the purpose of type checking.

Links to TS playground:

• Good
• Bad

Answer 1

An extremely simple solution would be to make a "dictionary of types", instead of unions and conditional types:

// Dictionary of types
type ShapeByType = {
    circle: Circle,
    square: Square,
    figure: Figure
}

function createShape<Args extends ShapeArgs>(args: Args): ShapeByType[Args["type"]] {
    switch (args.type) {
        case "circle": return new Circle(args);
        case "square": return new Square(args);
        case "figure": return new Figure(args);
    }
}

const circle1 = createShape({type: "circle", radius: 1});
//    ^? Circle

// ...

const circle2 = createShapeIn({type: "circle", radius: 1}, container);
//    ^? Circle

With this, it is easier to scale to an arbitrary number of Shapes.

Playground Link