TypeScript Support in Kea (The Long Road)

Even before Kea reached 1.0 last year, one topic kept popping up over and over again:

"Yeah it's great, but what about typescript?"

... or more eloquently:

"Unless the API has changed dramatically in the last few months it’s written in a way that ensure that it’s basically impossible to create effective typescript types to use it safely."

While that comment above is still technically true, as of version 2.2 (2.2.0-rc.1), Kea has full support for TypeScript!

The road there was long and winding... with plenty of dragons along the way.

Yet we prevailed!

But how?

What is Kea?

Kea is a state management library for React. Powered by Redux. It's like Redux Toolkit, but different, and older. It's designed to spark joy!

• Read "What is Kea?" to learn more.
• Open the "Quickstart" to see code.

TypeScript Support

First up, it's relatively easy to add TypeScript to a project. Just install the deps, convert your files to .ts or .tsx, set compilerOptions in tsconfig.json to strict and add types until there aren't any anys left.

This already gives a lot!

For example an autocomplete for resetContext:

But we want more. This should work as well:

and this:

How on earth do we do that?

The Gallery of Failed Attempts

As predicted by the Redditor quoted above:

"Unless the API has changed dramatically in the last few months it’s written in a way that ensure that it’s basically impossible to create effective typescript types to use it safely."

It turns out code like this is nearly impossible to type safely:

const logic = kea({
    // 1.
    actions: {
        openBlog: (id: number) => ({ id }), // 2.
    },
    reducers: (logic) => ({
        // 3.
        blog: [
            null,
            {
                openBlog: (_, { id }) => id, // 4.
                [appLogic.actions.closeBlog]: () => null, // 5.
            },
        ],
    }),
    selectors: {
        doubleBlog: [
            (s) => [s.blog], // 6.
            (blog) => (blog ? blog * 2 : null),
        ],
        tripleBlog: [
            (s) => [s.blog, s.doubleBlog], // 7.
            (blog, doubleBlog) => blog + doubleBlog,
        ],
    },
})

Who could have guessed?

There's a lot happening here:

1. We have a lot of keys (actions, reducers) inside one huge object literal {}
2. We have one action openBlog that takes an (id: number) and returns { id }
3. The reducers are specified as a function that gets the logic itself as its first parameter. That's some TS-killing loopy stuff right there!
4. The reducer blog uses the openBlog action (defined above in the same object!) to change its value
5. This reducer also depends on an action from a different logic, appLogic
6. The selector doubleBlog depends on the return type of the blog reducer
7. The selector tripleBlog depends on both blog and doubleBlog and their return types.

These are just a few of the complications.

This was going to be hard.

Yet I was determined to succeed, for I had on my side the strongest motivation on the planet: I had to prove someone wrong on the internet.

Attempt 1

It immediately became clear that just getting rid of anys in the codebase wasn't going to be enough.

The JavaScript that converts kea(input) into a logic is just a bit too complicated for the TypeScript compiler to automatically infer types from it.

TypeScript Generics enter the game.

Just write a long TypeScript type that gets the kea(input) parameter's type, looks at its properties and morphs them into a LogicType. Write some functional loopy stuff in a funny markup language. No big deal.

So I thought.

The first attempt looked like this when stripped to its core:

type Input = {
    actions?: (logic: Logic<Input>) => any // !
    reducers?: (logic: Logic<Input>) => any // !
    // ...
}
type Logic<I extends Input> = {
    actions: MakeLogicActions<I['actions']>
    reducers: MakeLogicReducers<I['reducers']>
    // ...
}
function kea<I extends Input>(input: I): Logic<I> {
    return realKea(input)
}

// helpers
type MakeLogicActions<InputActions> = {
    [K in keyof InputActions]: (
        ...args: Parameters<InputActions[K]>
    ) => {
        type: string
        payload: ReturnType<InputActions[K]>
    }
}
type MakeLogicReducers<InputReducers> = {
    // skip
}

This implementation gives us type completion when using the logic:

... but not when writing it:

The lines marked // ! are where this breaks down.

There's just no way to make the (logic: Logic<Input>) => any inside Input depend on the I extends Input that was passed to Logic<Input>.

Got all that? Me neither.

This kind of loopy stuff is just not possible with TypeScript:

// don't try this at home
type Input<L extends Logic<Input> = Logic<Input>> = {
    actions?: (logic: L) => MakeInputActions[Input['actions']] // ???
    reducers?: (logic: L) => MakeInputReducers[Input['actions']] // ???
    // ...
}
type Logic<I extends Input<Logic> = Input<Logic>> = {
    actions: MakeLogicActions<I['actions']>
    reducers: MakeLogicReducers<I['reducers']>
    // ...
}
function kea<I extends Input<Logic<I>>>(input: I): Logic<I> {
    return realKea(input)
}

With this attempt I got something to work, but ultimately without typing assistance inside the logic, it wouldn't prove someone on the internet wrong enough.

Back to the drawing board!

Attempt 2

I first alluded to automatic type generation 10 months ago, yet it always seemed like a huge undertaking. There had to be an easier way.

What if I changed the syntax of Kea itself to something friendlier to TypeScript? Hopefully in a completely opt-in and 100% backwards-compatible way?

Surely there won't be any problems maintaining two parallel implementations and everyone using Kea will understand that this is Kea's hooks moment [1], right? Right?

Right?

Would something like this be easier to type?

// pseudocode!
const logic = typedKea()
    .actions({
        submit: (id) => ({ id }),
        change: (id) => ({ id }),
    })
    .reducers({
        isSubmitting: [false, { submit: () => true }],
    })
    .listeners(({ actions }) => ({
        // (logic) => ...
        submit: async ({ id }) => {
            actions.change(id)
        },
    }))

I mean, it's just a slight alteration to this code that already works:

// real code!
const logic = kea({})
    .extend({
        actions: {
            submit: (id) => ({ id }),
            change: (id) => ({ id }),
        },
    })
    .extend({
        reducers: {
            isSubmitting: [false, { submit: () => true }],
        },
    })
    .extend({
        listeners: ({ actions }) => ({
            submit: async ({ id }) => {
                actions.change(id)
            },
        }),
    })

Surely not a big effort to refactor?

Unfortunately (or fortunately?), this approach didn't work either.

While this huge chain of type extensions sounds good in theory, you'll hit TypeScript's max instantiation depth limit eventually, as discovered by someone who was trying to add TS support to ... ehm, SQL?

I would experience the same. After a certain complexity the types just stopped working.

Definitely not ideal... and again won't prove someone on the internet wrong enough.

Attempt 3

Attempt 3 was one more go at attempt 1, but by building out the types in the other direction:

So instead of:

type Input = {
    actions?: (logic: Logic<Input>) => any
    reducers?: (logic: Logic<Input>) => any
}
type Logic<I extends Input> = {
    actions: MakeLogicActions<I['actions']>
    reducers: MakeLogicReducers<I['reducers']>
}
function kea<I extends Input>(input: I): Logic<I> {
    return realKea(input)
}

I started with something like:

interface AnyInput {}
export interface Input<A extends InputActions, R extends InputReducers, L extends InputListeners> extends AnyInput {
    actions?: InputActions
    reducers?: InputReducers<A, ReturnType<R>>
    listeners?: () => InputListeners<A, ReturnType<R>>
}

export interface Logic<I extends AnyInput> {
    /* This is a problem for another day! */
}

export declare function kea<T extends Input<T['actions'], T['reducers'], T['listeners']>>(input: T): Logic<T>

... only to fail even harder.

Attempt N+1

There were many other experiments and types that I tried.

They all had their issues.

In the end, it appears that this kind of loopy syntax that Kea uses together with selectors that depend on each other just wouldn't work with TypeScript.

That's even before you take into account plugins and logic.extend(moreInput).

What now?

I guess there's only one thing left to do.

My job now is to spend countless nights and weekends building kea-typegen, which will use the TypeScript Compiler API to load your project, analyse the generated AST, infer the correct types and write them back to disk in the form of logicType.ts files.

These logicTypes will then be fed back to the const logic = kea<logicType>() calls... and presto! Fully typed logic!

It's not ideal (ugh, another command to run), but it should work.

The stakes are high: If I fail or quit, the person on the internet will be proven right... and that is just not an option.

Automatic Type Generation

Thus it's with great excitement that I can announce kea-typegen to the world!

It's still rough with a lot of things to improve, yet it's already really useful!

We've been using it in PostHog for about a week now, and it's working great!

Take that, random person on the internet!

Install the typescript and kea-typegen packages, run kea-typegen watch and code away!

Read the TypeScript guide for more details.

Rough Edges

This is the very first version of kea-typegen, so there are still some rough edges.

1. You must manually import the logicType and insert it into your logic. This will be done automatically in the future.

2. You must manually hook up all type dependencies by adding them on the logicType in logic.ts. Kea-TypeGen will then put the same list inside logicType. This will also be done automatically in the future.

3. When connecting logic together, you must use [otherLogic.actionTypes.doSomething] instead of [otherLogic.actions.doSomething]

4. Sometimes you might need to "Reload All Files" in your editor at times... or explicitly open logicType.ts to see the changes.

5. Plugins aren't supported yet. I've hardcoded a few of them (loaders, router, window-values) into the typegen library, yet that's not a long term solution.

6. logic.extend() doesn't work yet.

These are all solvable issues. Let me know which ones to prioritise!

Alternative: MakeLogicType<V, A, P>

At the end of the day, Kea's loopy syntax doesn't bode well with TypeScript and we are forced to make our own logicTypes and feed them to Kea.

However nothing says these types need to be explicitly made by kea-typegen. You could easily make them by hand. Follow the example and adapt as needed!

To help with the most common cases, Kea 2.2.0 comes with a special type:

import { MakeLogicType } from 'kea'

type MyLogicType = MakeLogicType<Values, Actions, Props>

Pass it a bunch of interfaces denoting your logic's values, actions and props... and you'll get a close-enough approximation of the generated logic.

interface Values {
    id: number
    created_at: string
    name: string
    pinned: boolean
}

interface Actions {
    setName: (name: string) => { name: string }
}

interface Props {
    id: number
}

type RandomLogicType = MakeLogicType<Values, Actions, Props>

const randomLogic = kea<RandomLogicType>({
    /* skipping for brevity */
})

The result is a fully typed experience:

You'll even get completion when coding the logic:

Thank you to the team at Elastic for inspiring this approach!

Closing words

TypeScript support for Kea is finally here!

Well, almost. You can already use it in Kea v2.2.0-rc.1. The final v2.2.0 is not far away.

I've been building kea-typegen in isolation until now. I'd love to hear what the wider community thinks of it. Is it useful? What's missing? How can I improve the developer ergonomics? Can it work in your toolchain? Should I send the created logicTypes to GPT-3, so it would code the rest of your app? And who ate all the bagels?

Just open an issue and let's chat!

Also check out the samples folder in the kea-typegen repository for a few random examples of generated logic.

Finally here's a 12min video where I add TypeScript support to PostHog (we're hiring!):

Footnotes

[1] Hooks Moment: A massive improvement in developer ergonomics at the cost of all old code becoming legacy overnight.

Questions & Answers