refactor: improve TS definitions for VisitInfo class

Author: nodkz

src/VisitInfo.ts

@@ -16,17 +16,17 @@ import {
 } from './directoryToAst';
 import { FieldConfig } from './typeDefs';
 
-interface VisitInfoData<TContext = any> {
-  node: AstDirNode | AstFileNode | AstRootTypeNode;
+interface VisitInfoData<TNode extends AstDirNode | AstFileNode | AstRootTypeNode, TContext = any> {
+  node: TNode;
   nodeParent: AstDirNode | AstRootTypeNode | AstRootNode;
   operation: RootTypeNames;
   fieldName: string;
   fieldPath: string[];
   schemaComposer: SchemaComposer<TContext>;
 }
 
-export class VisitInfo<TContext = any> {
-  node: AstDirNode | AstFileNode | AstRootTypeNode;
+export class VisitInfo<TNode extends AstDirNode | AstFileNode | AstRootTypeNode, TContext = any> {
+  node: TNode;
   /** Parent AST node from directoryToAst */
   nodeParent: AstDirNode | AstRootTypeNode | AstRootNode;
   /** Brunch of schema under which is working visitor. Can be: query, mutation, subscription */
@@ -38,7 +38,7 @@ export class VisitInfo<TContext = any> {
   /** Type registry */
   schemaComposer: SchemaComposer<TContext>;
 
-  constructor(data: VisitInfoData<TContext>) {
+  constructor(data: VisitInfoData<TNode, TContext>) {
     this.node = data.node;
     this.operation = data.operation;
     this.nodeParent = data.nodeParent;
@@ -121,14 +121,13 @@ export class VisitInfo<TContext = any> {
    * This is mutable object and is shared between all calls.
    */
   get fieldConfig(): FieldConfig {
-    if (this.node.kind === 'file') {
-      return this.node.code?.default as FieldConfig;
-    } else if (this.node.kind === 'dir' || this.node.kind === 'rootType') {
-      return this.node.namespaceConfig?.code?.default as FieldConfig;
+    const node = this.node;
+    if (node.kind === 'file') {
+      return node.code?.default as FieldConfig;
+    } else if (node.kind === 'dir' || this.node.kind === 'rootType') {
+      return node.namespaceConfig?.code?.default as FieldConfig;
     }
-    throw new Error(
-      `Cannot get fieldConfig. Node has some strange kind: ${(this.node as any).kind}`
-    );
+    throw new Error(`Cannot get fieldConfig. Node has some strange kind: ${node.kind}`);
   }
 
   /**

src/astVisitor.ts

@@ -17,10 +17,10 @@ export type VisitorEmptyResult =
   | typeof VISITOR_REMOVE_NODE
   | typeof VISITOR_SKIP_CHILDREN;
 
-export type VisitKindFn<NodeKind> = (
+export type VisitKindFn<TNode extends AstDirNode | AstFileNode | AstRootTypeNode> = (
   /** Info & helper functions from visitor during traversing AST tree */
-  info: VisitInfo
-) => VisitorEmptyResult | NodeKind;
+  info: VisitInfo<TNode>
+) => VisitorEmptyResult | TNode;
 
 /**
  * Functions for every type of AST nodes which will be called by visitor.
@@ -66,14 +66,17 @@ export function astVisitor(
   });
 }
 
-export function visitNode(info: VisitInfo, visitor: AstVisitor): void {
+export function visitNode(
+  info: VisitInfo<AstDirNode | AstFileNode | AstRootTypeNode>,
+  visitor: AstVisitor
+): void {
   let result: VisitorEmptyResult | AstDirNode | AstFileNode | AstRootTypeNode;
   if (info.node.kind === 'dir') {
-    if (visitor.DIR) result = visitor.DIR(info);
+    if (visitor.DIR) result = visitor.DIR(info as VisitInfo<AstDirNode>);
   } else if (info.node.kind === 'file') {
-    if (visitor.FILE) result = visitor.FILE(info);
+    if (visitor.FILE) result = visitor.FILE(info as VisitInfo<AstFileNode>);
   } else if (info.node.kind === 'rootType') {
-    if (visitor.ROOT_TYPE) result = visitor.ROOT_TYPE(info);
+    if (visitor.ROOT_TYPE) result = visitor.ROOT_TYPE(info as VisitInfo<AstRootTypeNode>);
   }
 
   if (result === VISITOR_REMOVE_NODE) {