Typescript: Dynamic assignment to a typed map

Question

I'm writing some abstract entity system for fun where I have entities with traits. Traits have some fields, including dynamic data field:

enum TraitId {
  Movable = 'Movable', Rotatable = 'Rotatable', Scalable = 'Scalable', Collidable = 'Collidable'
}

interface TraitDataMovable {
  x: number;
  y: number;
}

type TraitDataMap = {
  [TraitId.Movable]: TraitDataMovable
  , [TraitId.Rotatable]: number // angle
  , [TraitId.Scalable]: number // scale
  , [TraitId.Collidable]: boolean // collides or not
}

interface TraitData<ID extends TraitId> {
  id: ID;
  data: TraitDataMap[ID];
  disabled?: boolean;
}

type EntityTraits = {
  [TID in TraitId]: TraitData<TID>
}

class Entity {
  id: string;
  traits: Partial<EntityTraits> = {};
}

So far I achieved correct behavior with manual assignment:

const ent = new Entity();

ent.traits.Rotatable = {
  id: TraitId.Rotatable, // id can only be Rotatable
  data: 100 // data can only be number
};

ent.traits.Collidable = {
  id: TraitId.Collidable, // id can only be Collidable
  data: true // data can only be boolean
}

const hasCollision = ent.traits.Collidable.data; // correctly typed as boolean

And now I'm trying to write function that adds any trait to the entity:

function addTraitToEntity(entity: Entity, traitData: TraitData<TraitId>) {
  entity.traits[traitData.id] = traitData;
  // Type 'TraitData<TraitId>' is not assignable to type 'undefined'.
}

function addTraitToEntity2<TID extends TraitId>(entity: Entity, traitData: TraitData<TID>) {
  entity.traits[traitData.id] = traitData;
  // Type 'TraitData<TID>' is not assignable to type 'Partial<EntityTraits>[TID]'.
  // Type 'TraitData<TID>' is not assignable to type 'undefined'
}

They work with // @ts-ignore, however I'd like to get rid of it and do it right. And understand how to have such system typed correctly.

Link to the playground

Answer 1

The problem is the following. TS somehow loses information about TID when trying to resolve the type of traitData.id expression. Consequently, the only information it has at the moment is the constraint to ... extends TraitId. That's why TS resolves traitData.id expression to TraitId or TraitId.Movable | TraitId.Scalable | TraitId.Rotatable | TraitId.Collidable instead of TID.

---

The simplest way to achieve the goal is to use the modified variant of addTraitToEntity2 function with existing types:

function addTraitToEntity2<TID extends TraitId>(entity: Entity, traitData: TraitData<TID>) {
  (entity.traits[traitData.id] as any) = traitData;
}

Negative aspects:

• as any is not a good way to solve issues;

Or as @Linda Paiste suggested, but as a generic function. Because without type variable you would be able to pass specific id member with an incorrect data member and vice versa.

function addTraitToEntity2<TID extends TraitId>(entity: Entity, traitData: TraitData<TID>) {
    entity.traits = {
        ...entity.traits,
        [traitData.id]: traitData
    };
}

---

Meantime, I refactored your code:

TS Playground

Improvements:

• Less number of TS types;
• Entity was introduced as an abstract class which includes API for add, remove, retrieve traits;
• Traits become a member of Entity class;
• Entity class is now generic and you can customize which traits entity can have;

Answer 2

Here's a roundabout solution to just avoid the error. It will work if you replace the entire object entity.traits instead of assigning just one property, like this:

function addTraitToEntity(entity: Entity, traitData: TraitData<TraitId>) {
    entity.traits = {
        ...entity.traits,
        [traitData.id]: traitData
    }
}

I've been playing around with this and I haven't really got a satisfying answer as to why your method doesn't work, but I think part of the problem is with typescript understanding that the key traitData.id matches the value traitData.