Typescript: Infer correct type from string literal union applied on keyed object

Question

I am using MobX and have this repository that contains all entities in my app.

I want to have two functions

• addEntity: Takes an entity and updates/inserts it
• getEntity: Returns the entity if found otherwise undefined

But I also want them to be typesafe.

Let's say that I keep this collection:

@observable entities = {
   authors: {} as Record<string, Author>,
   comments: {} as Record<string, Comment>,
   posts: {} as Record<string, Post>
}

I want to use my function like the following:

// "authors" should be checked against "authors" | "posts" | "comments"
// by specifying "authors" statically, I want typescript to automatically refine the return type to be Author, otherwise Author | Posts | Comment
addEntity("authors", new Author(...)) 

// same here
getEntity("authors", id) 

I have tried many ways with the generics but I can't get it right. I have to add a generic type that I type cast to So my calls are like this:

getEntity<Author>("authors")
// but nothing prevents me from writing
getEntity<Author>("posts")

Is there a trick to make this work?

Answer 1

Not sure about the MobX part, but the typescript type is a pretty straight forward application of index type queries and keyof :

class Author { id!: string; a!: string }
class Comment { id!: string; c!: string }
class Post {id!: string; p!: string }
class Store {
    entities = {
        authors: {} as Record<string, Author>,
        comments: {} as Record<string, Comment>,
        posts: {} as Record<string, Post>
    }
}
let store = new Store();
function addEntity<K extends keyof Store['entities']>(type: K, value: Store['entities'][K][string]) {
    let collection = store.entities[type] as Record<string, typeof value>; // some type assertions required
    collection[value.id] = value;
}
function getEntity<K extends keyof Store['entities']>(type: K, id: string): Store['entities'][K][string] {

    let collection = store.entities[type] as Record<string, Store['entities'][K][string]>; // some type assertions required
    return collection[id];
}
addEntity("authors", new Author()) 
addEntity("authors", new Post())  // err

// same here
let a: Author = getEntity("authors", "1") 
let p : Post = getEntity("authors", "1") //err