How to infer return type from parameter type in a generic function?

Question

Here is some code with conditional type

class A {
    public a: number;
}

class B {
    public b: number;
}

type DataType = "a" | "b";

type TData<T extends DataType> =
    T extends "a" ? A :
    T extends "b" ? B :
    never;

Now I want to use conditional type as a link from function parameter to its return type. I tried to achieve this in different ways with no result:

function GetData<T extends DataType>(dataType: T): TData<T> {
    if (dataType == "a")
        return new A();
    else if (dataType == "b")
        return new B();
}

What is the proper syntax? Is it possible with TypeScript 2.8?

Update

There is already an opened issue on github that covers my example. So current answer is "No, but may be possible in future".

Answer 1

You can use function overloads here:

function GetData(dataType: "a"): A;
function GetData(dataType: "b"): B;
function GetData(dataType: DataType): A | B {
    if (dataType === "a")
        return new A();
    else if (dataType === "b")
        return new B();
}

const f = GetData('a');  // Inferred A
const g = GetData('b');  // Inferred B

Answer 2

The ternary operator really doesn't play nice with generic types:

type TriviallyA<T extends DataType> = 
      T extends any ? A : A;
function GetData<T extends 'a' = 'a'>(dataType: T): TriviallyA<T> {
    return new A(); // Error: Type 'A' is not assignable to type 'TriviallyA<T>'
}

However generics do play nice with attribute lookup, so you can define an interface to map strings to particular types, then you can use keyof and attribute lookup to act as TData:

interface DataTypeMapping {
    a: A;
    b: B;
}
type DataType = keyof DataTypeMapping;
type TData<T extends DataType> = DataTypeMapping[T];

function GetData<T extends DataType>(dataType: T): TData<T> {
    // now expected return type is A | B so this is valid!
    if (dataType === 'a') {
        return new A(); 
    } else if (dataType === 'b') {
        return new B();
    }
}