Search code examples
typescript

Type inference for autocomplete param's value of an overloaded method

I'm trying to add some autocomplete support for a param of generic method

 subscribeToTelemetry<T extends keyof TypeMap>(name: T): Stream<TypeMap[T]>;
 subscribeToTelemetry<T extends NonNullable<{}>>(name: EntityName<T>): Stream<T>;

In a nutshell, if no type is provided subscribeToTelemetry the return type of should be inferred based on the name and some some injected/autogenerated types available. Additionally, possible name suggestions should be made using the injected/autogenerated types available. For example, using the following type:

type TypeMap = {
    ranger1: RangerTelemetry,
    ranger2: RangerTelemetry,
    mavic2: DroneTelemetry
};

and the following code

// should suggest: ['mavic2' | 'ranger2' | 'ranger1'] since no type T is provided 
const myStream1 = r.subscribeToTelemetry('ranger1'); // myStream1: Stream<RangerTelemetry

The suggested value should be ['mavic2' | 'ranger2' | 'ranger1'] since no type T is provided, and myStream1 type will be Stream<RangerTelemetry>

On the contrary, if we provide a type to subscribeToTelemetry and considering the the injected/autogenerated type available

const entityNames = {
    RangerTelemetry: ['ranger1', 'ranger2', 'ranger3'],
    DroneTelemetry: ['mavic1', 'mavic2', 'mavic3'],
} as const;

and some other type trics (see playground link), using the following code

// should suggest only: ['mavic1' | 'mavic2' | 'mavic3'] since type T is provided as DroneTelemetry,
const myStream3 = r.subscribeToTelemetry<DroneTelemetry>('mavic1');

We should suggest: ['mavic1' | 'mavic2' | 'mavic3'] since the type DroneTelemetry was provided, but ['mavic1' | 'mavic2' | 'mavic3', 'ranger1', 'ranger2', 'ranger3'] is shown instead. The type of myStream3 will be Stream<DroneTelemetry>

I have multiple variants for the signature implementation of the overloaded subscribeToTelemetry method without succeeding to provide the right suggestions when providing a type, I always get the rangers suggestion even when I specify DroneTelemetry as type.

class Registry {
    subscribeToTelemetry<T extends keyof TypeMap>(name: T): Stream<TypeMap[T]>;
    subscribeToTelemetry<T extends NonNullable<{}>>(name: EntityName<T>): Stream<T>;
    subscribeToTelemetry<T extends NonNullable<{}>>(name: EntityName<T>): Stream<T> {
        return new Stream<T>(name); 
    }
}

If anyone has any suggestions, I will appreciate any help

Humberto

Here is the full code:

type DroneTelemetry = { kind: 'DroneTelemetry', batteryLevel: number, velocity: number, location: {latitude: number, longitude: number}}
type RangerTelemetry = { kind: 'RangerTelemetry', velocity: number, location: {latitude: number, longitude: number}}


////// the entity names and type mapping will be autogenerated
const entityNames = {
    RangerTelemetry: ['ranger1', 'ranger2', 'ranger3'],
    DroneTelemetry: ['mavic1', 'mavic2', 'mavic3'],
} as const;

type TypeMap = {
    ranger1: RangerTelemetry,
    ranger2: RangerTelemetry,
    mavic2: DroneTelemetry
};
////////

// Define the types for the keys and values of the mapping
type EntityNames = typeof entityNames;
type EntityKeys = keyof EntityNames;
type EntityValues<T extends EntityKeys> = EntityNames[T][number];

// Define a mapping from type aliases to their corresponding keys
type TypeToKey<T> = T extends { kind: infer K } ? K : never;

type ExtractEntitiesNames<T> = TypeToKey<T> extends keyof EntityNames ? EntityValues<TypeToKey<T>> : never;
type EntityName<T> = ExtractEntitiesNames<T> |  (string  & {});

/////////////////////////////////////////////////////////////////////

class Stream<T> {
    name: string
    constructor(name: string){
        this.name = name
    }
    getLastValue() {
        return {} as T // dummy impl
    }
}

class Registry {
    subscribeToTelemetry<T extends keyof TypeMap>(name: T): Stream<TypeMap[T]>;
    subscribeToTelemetry<T extends NonNullable<{}>>(name: EntityName<T>): Stream<T>;
    subscribeToTelemetry<T extends NonNullable<{}>>(name: EntityName<T>): Stream<T> {
        return new Stream<T>(name); 
    }
}

const r = new Registry();
// should suggest: ['mavic2' | 'ranger2' | 'ranger1'] since no type T is provided // OK
const myStream1 = r.subscribeToTelemetry('ranger1'); // myStream1 type will be Stream<RangerTelemetry>

// should suggest: ['mavic2' | 'ranger2' | 'ranger1'] since no type T is provided // OK
const myStream2 = r.subscribeToTelemetry('ranger2'); // myStream2 type will be Stream<RangerTelemetry>

 // should suggest only: ['mavic1' | 'mavic2' | 'mavic3'] since type T is provided as DroneTelemetry,
 // but ['mavic1' | 'mavic2' | 'mavic3', 'ranger1', 'ranger2', 'ranger3'] is shown instead.                // WRONG
const myStream3 = r.subscribeToTelemetry<DroneTelemetry>('mavic1'); // myStream3 type will be Stream<RangerTelemetry>

Here is a ts-playground with the full code

Solution

  • We can fix this by changing the approach to how we map the types to each other. This will be better seen in the code than explained. With this approach you don't need to overload your function.

    A summary of the approach:

    • Index the TypeMap (with Record<> or {[x in key]: value} where drone names have a value of DroneTelemetry and the same for the ranger names
    • Define two generic types for subscribeToTelemetry
      • T is the entity type DroneTelemetry|RangerTelemetry
      • F is the entity name based on T, which when not narrowed, either by definition or inference, will be all the entity names. When it's narrowed by a passed T, it's only the keys that pertain to the entity type thanks to EntityName. This isn't just useful for parameter autocomplete, but for the return type to, as per the next point.
    • Rely on passing the key F to the TypeMap to get the returned EntityType
      • If T isn't passed. F's value will be inferred, leading to the matched type
      • If the entity type T is passed, we don't want to pass F (too hectic), so we have a default type which narrows to only the names that pertain to that entity type, which means that we've narrowed the key we pass to TypeMap, returning the relevant EntityType

    One thing to note, the second generic type F is inferred when we don't pass T, but when we pass T, we have to pass F, because you can't pass some generic types and infer the rest, it's all or none. That's why we need to give F a default type.

    type DroneTelemetry = { kind: 'DroneTelemetry', batteryLevel: number, velocity: number, location: { latitude: number, longitude: number } }
type RangerTelemetry = { kind: 'RangerTelemetry', velocity: number, location: { latitude: number, longitude: number } }

type EntityType = DroneTelemetry | RangerTelemetry;

const entityNames = {
    RangerTelemetry: ['ranger1', 'ranger2', 'ranger3'],
    DroneTelemetry: ['mavic1', 'mavic2', 'mavic3'],
} as const;

type EntityNameMap = typeof entityNames;
type EntityName<T extends EntityType> = EntityNameMap[T['kind']][number]

type TypeMap = Record<EntityName<RangerTelemetry>, RangerTelemetry> &
    Record<EntityName<DroneTelemetry>, DroneTelemetry>

class Stream<T extends EntityType> {
    name: string
    constructor(name: EntityName<T>) {
        this.name = name
    }
    getLastValue() {
        return {} as T // dummy impl
    }
}

class Registry {
    subscribeToTelemetry<T extends EntityType, F extends EntityName<T> = EntityName<T>>(name: F): Stream<TypeMap[F]> {
        return new Stream<TypeMap[F]>(name);
    }
}

const r = new Registry();
// Name options are all entity names
const myStream1 = r.subscribeToTelemetry('mavic2');
//        ^?const myStream2: Stream<DroneTelemetry>
const myStream2 = r.subscribeToTelemetry('ranger2');
//       ^?const myStream2: Stream<RangerTelemetry>

// Options are only drone names
const myStream3 = r.subscribeToTelemetry<DroneTelemetry>('mavic2');
//       ^?const myStream3: Stream<DroneTelemetry>

    Playground