feat: add ChildTypes helper

HerringtonDarkholme · HerringtonDarkholme · commit 0d477e9fe8d6 · 2024-12-17T13:15:30.000-05:00
diff --git a/crates/napi/__test__/type.spec.ts b/crates/napi/__test__/type.spec.ts
@@ -17,7 +17,7 @@ test('test no type annotation', t => {
   t.assert(childKind !== ',')
   // test parent method
   const parent = child!.parent()
-  t.assert(parent === root)
+  t.assert(parent!.kind() === root.kind())
   // test find
   const sum = root.find({
     rule: {
@@ -65,7 +65,7 @@ test('test type assertion', t => {
   t.assert(childKind !== ',')
   // test parent method
   const parent = child!.parent()
-  t.assert(parent === root)
+  t.assert(parent?.kind() === root.kind())
   // test find
   const sum = root.find({
     rule: {
@@ -114,13 +114,15 @@ test('test type argument style', t => {
   t.assert(childKind !== ',')
   // test parent method
   const parent = child!.parent<'program'>()
-  t.assert(parent === root)
+  t.assert(parent!.kind() === root.kind())
   // @ts-expect-error should reject wrong kind
   const parent2 = child!.parent<'eskdf'>()
-  t.assert(parent2 === root)
+  t.assert(parent2!.kind() === root.kind())
   const parent3 = child!.parent<'expression_statement'>()
   // @ts-expect-error parent3 should take new kind
-  t.assert(parent3 === root)
+  t.assert(parent3!.kind() === root.kind())
+  // @ts-expect-error should report invalid child type
+  root.child<'binary_expression'>
   // test find
   const sum = root.find<'binary_expression'>({
     rule: {
diff --git a/crates/napi/types/sgnode.d.ts b/crates/napi/types/sgnode.d.ts
@@ -7,6 +7,8 @@ import type {
   NodeFieldInfo,
   RootKind,
   NamedKinds,
+  ChildKinds,
+  NamedChildKinds,
 } from './staticTypes'
 import type { NapiConfig } from './config'
 
@@ -71,7 +73,8 @@ export declare class SgNode<
     name: F,
   ): Exclude<FieldNode<M, T, F>, null>[]
   parent: NodeMethod<M>
-  child<K extends Kinds<M>>(nth: number): RefineNode<M, K> | null
+  child(nth: number): SgNode<M, ChildKinds<M, T>> | null
+  child<K extends NamedChildKinds<M, T>>(nth: number): RefineNode<M, K> | null
   ancestors(): Array<SgNode<M>>
   next: NodeMethod<M>
   nextAll(): Array<SgNode<M>>
@@ -91,9 +94,9 @@ export declare class SgRoot<M extends TypesMap = TypesMap> {
   filename(): string
 }
 
-interface NodeMethod<M extends TypesMap, Args extends unknown[] = [], T = NamedKinds<M>> {
+interface NodeMethod<M extends TypesMap, Args extends unknown[] = []> {
   (...args: Args): SgNode<M> | null
-  <K extends T>(...args: Args): RefineNode<M, K> | null
+  <K extends NamedKinds<M>>(...args: Args): RefineNode<M, K> | null
 }
 
 /**
diff --git a/crates/napi/types/staticTypes.d.ts b/crates/napi/types/staticTypes.d.ts
@@ -48,12 +48,21 @@ export type ExtractField<
   : NodeFieldInfo
 
 // in case of empty types array, return string as fallback
-type NoNever<T, Fallback = string> = [T] extends [never] ? Fallback : T
+type NoNever<T, Fallback> = [T] extends [never] ? Fallback : T
 
 export type TypesInField<M extends TypesMap, I extends NodeFieldInfo> = NoNever<
-  ResolveType<M, I['types'][number]['type']>
+  ResolveType<M, I['types'][number]['type']>,
+  Kinds<M>
 >
 
+export type NamedChildKinds<M extends TypesMap, T extends Kinds<M>> =
+  M[T] extends { children: infer C extends NodeFieldInfo }
+    ? TypesInField<M, C>
+    : NamedKinds<M>
+export type ChildKinds<M extends TypesMap, T extends Kinds<M>> =
+  | NamedChildKinds<M, T>
+  | LowPriorityKey
+
 /**
  * resolve subtypes alias. see tree-sitter's reference
  * e.g. like `expression` => `binary_expression` | `unary_expression` | ...
