From d45ad500e75b2617ad12fdebb72ad40e13486264 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=8E=8B=E6=B8=85=E9=9B=A8?= <colinwang.0616@gmail.com>
Date: Thu, 29 Jul 2021 03:47:29 +0800
Subject: [PATCH] Improve `path.get` typings with TS template literal types
 (#13588)

---
 packages/babel-traverse/src/path/family.ts | 44 ++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/packages/babel-traverse/src/path/family.ts b/packages/babel-traverse/src/path/family.ts
index 8370195236..4bf8736900 100644
--- a/packages/babel-traverse/src/path/family.ts
+++ b/packages/babel-traverse/src/path/family.ts
@@ -320,6 +320,44 @@ export function getAllPrevSiblings(this: NodePath): NodePath[] {
   return siblings;
 }
 
+// convert "1" to 1 (string index to number index)
+type MaybeToIndex<T extends string> = T extends `${bigint}` ? number : T;
+
+type Pattern<Obj extends string, Prop extends string> = `${Obj}.${Prop}`;
+
+// split "body.body.1" to ["body", "body", 1]
+type Split<P extends string> = P extends Pattern<infer O, infer U>
+  ? [MaybeToIndex<O>, ...Split<U>]
+  : [MaybeToIndex<P>];
+
+// get all K with Node[K] is t.Node | t.Node[]
+type NodeKeyOf<Node extends t.Node | t.Node[]> = keyof Pick<
+  Node,
+  {
+    [Key in keyof Node]-?: Node[Key] extends t.Node | t.Node[] ? Key : never;
+  }[keyof Node]
+>;
+
+// traverse the Node with tuple path ["body", "body", 1]
+// Path should be created with Split
+type Trav<
+  Node extends t.Node | t.Node[],
+  Path extends unknown[],
+> = Path extends [infer K, ...infer R]
+  ? K extends NodeKeyOf<Node>
+    ? R extends []
+      ? Node[K]
+      : // @ts-expect-error ignore since TS is not smart enough
+        Trav<Node[K], R>
+    : never
+  : never;
+
+type ToNodePath<T> = T extends Array<t.Node | null | undefined>
+  ? Array<NodePath<T[number]>>
+  : T extends t.Node | null | undefined
+  ? NodePath<T>
+  : never;
+
 function get<T extends t.Node, K extends keyof T>(
   this: NodePath<T>,
   key: K,
@@ -330,6 +368,12 @@ function get<T extends t.Node, K extends keyof T>(
   ? NodePath<T[K]>
   : never;
 
+function get<T extends t.Node, K extends string>(
+  this: NodePath<T>,
+  key: K,
+  context?: boolean | TraversalContext,
+): ToNodePath<Trav<T, Split<K>>>;
+
 function get<T extends t.Node>(
   this: NodePath<T>,
   key: string,