Implementing Ruby's `.dig` Hash method in TypeScript

20 Apr, 2023

What's .dig?

Ruby hashes have a .dig method that allows us safely access a nested value in a Hash.

h = { foo: { bar: { baz: 12 } } }

puts h.dig(:foo, :bar, :baz) # 12
puts h.dig(:foo, :qux, :baz) # nil (no Error)

h?.foo?.bar?.baz; // 12
h?.foo?.qux?.baz; // undefined (no Error)

h.get('foo').get('bar').get('baz') # 12
h.get('foo', {}).get('qux', {}).get('baz') # None (no Error if we provide the second arguments)

It practically writes itself! A simplified, but fully-functioning implementation.

The initial implementation of it is actually pretty intuitive. We want a generic type that takes in two slots: the "object" and the "path":

type Dig<Obj, Path> = { /* ... */ };

Then we want to ensure that the first element of Path is a keyof Obj. To get the first element, we'll need to use infer along with TypeScript's awkward but powerful extends ternaries to do some conditional logic:

type Dig<Obj, Path> =
  Path extends [infer FirstKey, ...infer Rest] // infer our `FirstKey` and `Rest`
    ? FirstKey extends keyof Obj // if FirstKey is a valid key,
      ? Obj[FirstKey]            // return the accessed value;
      : Obj                      // otherwise return the original object
    : never
;

Great we've actually already got our first key working, we just need to make this recursive. To do that we'll call our Dig generic again but this time pass our Obj[FirstKey] in in the Obj slot and Rest in the Path slot:

type Dig<Obj, Path> =
  Path extends [infer FirstKey, ...infer Rest]
    ? FirstKey extends keyof Obj
      ? Dig<Obj[FirstKey], Rest>
      : Obj
    : Obj
;

You'll noticed I changed the last line from never to passing back the original Obj. This is because we have no constraints on our Path slot so we're always doing a recursion even after we're out of elements.

And it works!

type What1 = Dig<{ foo: { bar: { baz: string } } }, ['foo', 'bar', 'baz']>;
//    ^? string
type What2 = Dig<{ foo: { bar: { baz: string } } }, ['foo', 'bar', 'nonsense']>;
//    ^? { baz: string; }

Making TypeScript mad on purpose Adding type safety to our Path generic slot.

You're probably thinking, "Wait, isn't the point of TypeScript to tell us an issue before it happens? What if we want to use logic like this to warn us if we give it a bad Path?"

To do this, we will add a generic constraint to the Path slot. There's a couple of approaches we could take (recursive types do work in generic slot constraints) but perhaps the most straightforward is to simply use a utility type that will get us all possible valid paths:

type PossiblePath<T> = any;

type Dig<
  Obj, 
  Path extends PossiblePath<Obj>
> =
  Path extends [infer FirstKey, ...infer Rest]
    ? FirstKey extends keyof Obj
      ? Rest extends PossiblePath<Obj[FirstKey]>
        ? Dig<Obj[FirstKey], Rest>
        : Obj[FirstKey]
      : Obj
    : Obj
;

As you can see, adding this constraint on the Path slot required us to add an explicit check to assure TypeScript we know what we're doing.

So now at least our job is rather straightforward. We just need to write a PossiblePath type that takes some object and returns a union of all the possible valid paths.

Let's take this as a reasonable first guess:

type Primitive = string | number | boolean;

type PossiblePath<T> = [
  keyof T, // first element is just the keys of T
  // since we have a tuple of a variable length we'll use ... notation
  ...(T[keyof T] extends Primitive
    ? [] // if the value is a primitive, we don't need to get the `keyof` of it
    : PossiblePath<T[keyof T]> // if, instead, it's an object we recurse
  )
];

And let's concoct an example to see how we're looking.

const example = {
  foo: { bar: 'baz', qux: { quux: 3 } },
  corge: { bar: true }
};

declare const blah: PossiblePath<typeof example>;
//             ^? ["foo" | "corge", "bar"]

It looks like that first element is pretty close, we just want it to be a union of tuples rather than a tuple of unions. That is to say, we want ["foo", "bar"] | ["corge", "bar"] rather than ["foo" | "corge", "bar"].

So that's one issue. The second (similar) issue is that our second element is keyof Obj["foo" | "bar"] and we instead want to have keyof Obj["foo"] | keyof Obj["corge"].

The issue we're trying to solve here is that our unions are not discriminated. We basically want to "loop" through a TypeScript union. One ingenious (hack?) I found to accomplish this is to wrap our definition in a conditional check.

type PossiblePath<
  T, 
  Key extends keyof T = keyof T // introduce Key as a generic
> = Key extends Key ? [ // use a conditional to discriminate 
  Key,
  ...(T[Key] extends Primitive ? [] : PossiblePath<T[Key]>)
] : never;

So we added our Key as a generic slot and gave it the same default value as before (keyof T) so that we don't have to do any extra work. Then we used this silly Key extends Key bit to wrap our previous logic in a conditional type. This forces TypeScript to check each member of the string union. The result? Revisiting our previous example, we now have:

declare const blah: PossiblePath<typeof example>;
//             ^? ["foo", "bar"] | ["foo", "qux", "quux"] | ["corge", "bar"]

This looks great! Just one final small issue. Currently our result is only the leaf paths. But we want all the incomplete paths from along the way too! One final hack is to turn ? [] : PossiblePath<T[Key]>) into ? [] : PossiblePath<T[Key]> | []). This forces the logic to add "short-circuited" paths to our tuple union.

And we're done! Our final code looks like:

// #region motivating example
const example = {
  foo: { bar: 'baz', qux: { quux: 3 } },
  corge: { bar: true }
};

// #region PossiblePath
type Primitive = string | number | boolean;

type PossiblePath<T, Key extends keyof T = keyof T> = Key extends Key ? [
  Key,
  ...(T[Key] extends Primitive ? [] : PossiblePath<T[Key]> | [])
] | [] : never;

declare const test: PossiblePath<typeof example>;
//             ^? [] | ["foo"] | ["foo", "bar"] | ["foo", "qux"] | ["foo", "qux", "quux"] | ["corge"] | ["corge", "bar"]

// #region Dig
type Dig<Obj, Path extends PossiblePath<Obj>> =
  Path extends [infer FirstKey, ...infer Rest]
    ? FirstKey extends keyof Obj
      ? Rest extends PossiblePath<Obj[FirstKey]>
        ? Dig<Obj[FirstKey], Rest>
        : Obj[FirstKey]
      : Obj
    : Obj
;

type ShouldSucceed = Dig<typeof example, ['foo', 'qux', 'quux']>
//    ^? number
type ShouldError = Dig<typeof example, ['foo', 'nonsense', 'quux']>
//    ^? ERROR: Type '"nonsense"' is not assignable to type '"qux"'

Try it out for yourself in the TypeScript Playground.

#dictionary-methods #ruby #typescript