feat(react): improve lib schematic by replacing parentRoute with appProject option

- Instead of providing full path to app component, just pass the project's name.
- Updates both app.tsx and main.tsx so user no longer need to do any maual updates.

docs/api-react/schematics/library.md

@@ -11,6 +11,12 @@ ng generate library ...
 
 ## Options
 
+### appProject
+
+Type: `string`
+
+The application project to add the library route to
+
 ### directory
 
 Type: `string`
@@ -31,12 +37,6 @@ Type: `string`
 
 Library name
 
-### parentRoute
-
-Type: `string`
-
-Add new route to the parent component as specified by this path
-
 ### pascalCaseFiles
 
 Default: `false`

packages/react/src/schematics/application/application.ts

@@ -8,6 +8,7 @@ import {
   move,
   noop,
   Rule,
+  SchematicContext,
   template,
   Tree,
   url
@@ -35,7 +36,7 @@ import * as ts from 'typescript';
 
 import { Schema } from './schema';
 import { CSS_IN_JS_DEPENDENCIES } from '../../utils/styled';
-import { addRouter } from '../../utils/ast-utils';
+import { addRoute, addInitialRoutes } from '../../utils/ast-utils';
 import {
   babelCoreVersion,
   babelLoaderVersion,
@@ -61,7 +62,7 @@ interface NormalizedSchema extends Schema {
 }
 
 export default function(schema: Schema): Rule {
-  return (host: Tree) => {
+  return (host: Tree, context: SchematicContext) => {
     const options = normalizeOptions(host, schema);
 
     return chain([
@@ -89,7 +90,7 @@ export default function(schema: Schema): Rule {
           })
         : noop(),
       addStyledModuleDependencies(options),
-      addRouting(options),
+      addRouting(options, context),
       addBabel(options),
       formatFiles(options)
     ]);
@@ -221,7 +222,10 @@ function addStyledModuleDependencies(options: NormalizedSchema): Rule {
     : noop();
 }
 
-function addRouting(options: NormalizedSchema): Rule {
+function addRouting(
+  options: NormalizedSchema,
+  context: SchematicContext
+): Rule {
   return options.routing
     ? chain([
         function addRouterToComponent(host: Tree) {
@@ -237,7 +241,7 @@ function addRouting(options: NormalizedSchema): Rule {
             true
           );
 
-          insert(host, appPath, addRouter(appPath, appSource));
+          insert(host, appPath, addInitialRoutes(appPath, appSource, context));
         },
         addDepsToPackageJson({ 'react-router-dom': reactRouterVersion }, {})
       ])

packages/react/src/schematics/library/library.spec.ts

@@ -263,8 +263,8 @@ describe('lib', () => {
     });
   });
 
-  describe('--parentRoute', () => {
-    it('should add route to parent component', async () => {
+  describe('--appProject', () => {
+    it('should add new route to existing routing code', async () => {
       appTree = await runSchematic(
         'app',
         { name: 'myApp', routing: true },
@@ -275,48 +275,39 @@ describe('lib', () => {
         'lib',
         {
           name: 'myLib',
-          parentRoute: 'apps/my-app/src/app/app.tsx'
+          appProject: 'my-app'
         },
         appTree
       );
 
       const appSource = tree.read('apps/my-app/src/app/app.tsx').toString();
+      const mainSource = tree.read('apps/my-app/src/main.tsx').toString();
 
-      expect(appSource).toMatch(/<Route\s*path="\/my-lib"/);
+      expect(mainSource).toContain('react-router-dom');
+      expect(mainSource).toContain('<BrowserRouter>');
       expect(appSource).toContain('@proj/my-lib');
+      expect(appSource).toContain('react-router-dom');
+      expect(appSource).toMatch(/<Route\s*path="\/my-lib"/);
     });
 
-    it('throws error when parent component file is missing', async () => {
-      await expect(
-        runSchematic(
-          'lib',
-          {
-            name: 'myLib',
-            parentRoute: 'does/not/exist.tsx'
-          },
-          appTree
-        )
-      ).rejects.toThrow('Cannot find');
-    });
-
-    it('should add routing to app if it does not exist yet', async () => {
-      appTree = await runSchematic(
-        'app',
-        { name: 'myApp', routing: false },
-        appTree
-      );
+    it('should initialize routes if none were set up then add new route', async () => {
+      appTree = await runSchematic('app', { name: 'myApp' }, appTree);
 
       const tree = await runSchematic(
         'lib',
         {
           name: 'myLib',
-          parentRoute: 'apps/my-app/src/app/app.tsx'
+          appProject: 'my-app'
         },
         appTree
       );
 
       const appSource = tree.read('apps/my-app/src/app/app.tsx').toString();
+      const mainSource = tree.read('apps/my-app/src/main.tsx').toString();
 
+      expect(mainSource).toContain('react-router-dom');
+      expect(mainSource).toContain('<BrowserRouter>');
+      expect(appSource).toContain('@proj/my-lib');
       expect(appSource).toContain('react-router-dom');
       expect(appSource).toMatch(/<Route\s*path="\/my-lib"/);
     });

packages/react/src/schematics/library/library.ts

@@ -13,8 +13,11 @@ import {
 } from '@angular-devkit/schematics';
 import {
   addDepsToPackageJson,
+  addLintFiles,
   formatFiles,
+  generateProjectLint,
   getNpmScope,
+  getProjectConfig,
   insert,
   names,
   NxJson,
@@ -23,28 +26,34 @@ import {
   toClassName,
   toFileName,
   updateJsonInTree,
-  updateWorkspaceInTree,
-  addLintFiles,
-  generateProjectLint
+  updateWorkspaceInTree
 } from '@nrwl/workspace';
 import { join, normalize, Path } from '@angular-devkit/core';
 import * as ts from 'typescript';
 
 import { Schema } from './schema';
-import { addRoute, addRouter } from '../../utils/ast-utils';
+import {
+  addBrowserRouter,
+  addInitialRoutes,
+  addRoute,
+  findComponentImportPath
+} from '../../utils/ast-utils';
 import { reactRouterVersion } from '../../utils/versions';
 
 export interface NormalizedSchema extends Schema {
   name: string;
   fileName: string;
   projectRoot: Path;
+  routePath: string;
   projectDirectory: string;
   parsedTags: string[];
+  appMain?: string;
+  appSourceRoot?: Path;
 }
 
 export default function(schema: Schema): Rule {
   return (host: Tree, context: SchematicContext) => {
-    const options = normalizeOptions(schema);
+    const options = normalizeOptions(host, schema, context);
 
     return chain([
       addLintFiles(options.projectRoot, options.linter),
@@ -68,7 +77,7 @@ export default function(schema: Schema): Rule {
         export: true,
         routing: options.routing
       }),
-      updateParentRoute(options),
+      updateAppRoutes(options, context),
       formatFiles(options)
     ])(host, context);
   };
@@ -132,45 +141,77 @@ function updateNxJson(options: NormalizedSchema): Rule {
   });
 }
 
-function updateParentRoute(options: NormalizedSchema): Rule {
-  return options.parentRoute
-    ? chain([
-        function ensureRouterAdded(host: Tree) {
-          const { source, content } = readComponent(host, options.parentRoute);
-          const isRouterPresent = content.match(/react-router-dom/);
+function updateAppRoutes(
+  options: NormalizedSchema,
+  context: SchematicContext
+): Rule {
+  if (!options.appMain || !options.appSourceRoot) {
+    return noop();
+  }
+  return (host: Tree) => {
+    const { source } = readComponent(host, options.appMain);
+    const componentImportPath = findComponentImportPath('App', source);
 
-          if (!isRouterPresent) {
-            insert(
-              host,
-              options.parentRoute,
-              addRouter(options.parentRoute, source)
-            );
-            return addDepsToPackageJson(
-              { 'react-router-dom': reactRouterVersion },
-              {}
-            );
-          }
-        },
-        function addRouteToComponent(host: Tree) {
-          const npmScope = getNpmScope(host);
-          const { source: componentSource } = readComponent(
-            host,
-            options.parentRoute
-          );
+    if (!componentImportPath) {
+      throw new Error(
+        `Could not find App component in ${
+          options.appMain
+        } (Hint: you can omit --appProject, or make sure App exists)`
+      );
+    }
 
+    const appComponentPath = join(
+      options.appSourceRoot,
+      `${componentImportPath}.tsx`
+    );
+    return chain([
+      addDepsToPackageJson({ 'react-router-dom': reactRouterVersion }, {}),
+      function addBrowserRouterToMain(host: Tree) {
+        const { content, source } = readComponent(host, options.appMain);
+        const isRouterPresent = content.match(/react-router-dom/);
+        if (!isRouterPresent) {
           insert(
             host,
-            options.parentRoute,
-            addRoute(options.parentRoute, componentSource, {
-              libName: options.name,
+            options.appMain,
+            addBrowserRouter(options.appMain, source, context)
+          );
+        }
+      },
+      function addInitialAppRoutes(host: Tree) {
+        const { content, source } = readComponent(host, appComponentPath);
+        const isRouterPresent = content.match(/react-router-dom/);
+        if (!isRouterPresent) {
+          insert(
+            host,
+            appComponentPath,
+            addInitialRoutes(appComponentPath, source, context)
+          );
+        }
+      },
+      function addNewAppRoute(host: Tree) {
+        const npmScope = getNpmScope(host);
+        const { source: componentSource } = readComponent(
+          host,
+          appComponentPath
+        );
+        insert(
+          host,
+          appComponentPath,
+          addRoute(
+            appComponentPath,
+            componentSource,
+            {
+              routePath: options.routePath,
               componentName: toClassName(options.name),
               moduleName: `@${npmScope}/${options.projectDirectory}`
-            })
-          );
-        },
-        addDepsToPackageJson({ 'react-router-dom': reactRouterVersion }, {})
-      ])
-    : noop();
+            },
+            context
+          )
+        );
+      },
+      addDepsToPackageJson({ 'react-router-dom': reactRouterVersion }, {})
+    ]);
+  };
 }
 
 function readComponent(
@@ -193,26 +234,54 @@ function readComponent(
   return { content, source };
 }
 
-function normalizeOptions(options: Schema): NormalizedSchema {
+function normalizeOptions(
+  host: Tree,
+  options: Schema,
+  context: SchematicContext
+): NormalizedSchema {
   const name = toFileName(options.name);
   const projectDirectory = options.directory
     ? `${toFileName(options.directory)}/${name}`
     : name;
 
   const projectName = projectDirectory.replace(new RegExp('/', 'g'), '-');
-  const fileName = options.simpleModuleName ? name : projectName;
+  const fileName = projectName;
   const projectRoot = normalize(`libs/${projectDirectory}`);
 
   const parsedTags = options.tags
     ? options.tags.split(',').map(s => s.trim())
     : [];
 
-  return {
+  const normalized: NormalizedSchema = {
     ...options,
     fileName,
+    routePath: `/${name}`,
     name: projectName,
     projectRoot,
     projectDirectory,
     parsedTags
   };
+
+  if (options.appProject) {
+    const appProjectConfig = getProjectConfig(host, options.appProject);
+
+    if (appProjectConfig.projectType !== 'application') {
+      throw new Error(
+        `appProject expected type of "application" but got "${
+          appProjectConfig.projectType
+        }"`
+      );
+    }
+
+    try {
+      normalized.appMain = appProjectConfig.architect.build.options.main;
+      normalized.appSourceRoot = normalize(appProjectConfig.sourceRoot);
+    } catch (e) {
+      throw new Error(
+        `Could not locate project main for ${options.appProject}`
+      );
+    }
+  }
+
+  return normalized;
 }

packages/react/src/schematics/library/schema.d.ts

@@ -7,10 +7,9 @@ export interface Schema {
   skipTsConfig: boolean;
   skipFormat: boolean;
   tags?: string;
-  simpleModuleName: boolean;
   pascalCaseFiles?: boolean;
   routing?: boolean;
-  parentRoute?: string;
+  appProject?: string;
   unitTestRunner: 'jest' | 'none';
   linter: Linter;
 }

packages/react/src/schematics/library/schema.json

@@ -84,9 +84,9 @@
       "type": "boolean",
       "description": "Generate library with routes"
     },
-    "parentRoute": {
+    "appProject": {
       "type": "string",
-      "description": "Add new route to the parent component as specified by this path"
+      "description": "The application project to add the library route to"
     }
   },
   "required": ["name"]

packages/react/src/utils/ast-utils.spec.ts

@@ -1,16 +1,18 @@
 import * as utils from './ast-utils';
 import * as ts from 'typescript';
+import { Tree } from '@angular-devkit/schematics';
+import { insert } from '@nrwl/workspace/src/utils/ast-utils';
 
 describe('react ast-utils', () => {
   describe('findDefaultExport', () => {
     it('should find exported variable', () => {
-      const text = `
+      const sourceCode = `
         const main = () => {}; 
         export default main;
       `;
       const source = ts.createSourceFile(
         'test.ts',
-        text,
+        sourceCode,
         ts.ScriptTarget.Latest,
         true
       );
@@ -22,13 +24,13 @@ describe('react ast-utils', () => {
     });
 
     it('should find exported function', () => {
-      const text = `
+      const sourceCode = `
         function main() {} 
         export default main;
       `;
       const source = ts.createSourceFile(
         'test.ts',
-        text,
+        sourceCode,
         ts.ScriptTarget.Latest,
         true
       );
@@ -40,12 +42,12 @@ describe('react ast-utils', () => {
     });
 
     it('should find default export function', () => {
-      const text = `
+      const sourceCode = `
         export default function main() {};
       `;
       const source = ts.createSourceFile(
         'test.ts',
-        text,
+        sourceCode,
         ts.ScriptTarget.Latest,
         true
       );
@@ -56,4 +58,133 @@ describe('react ast-utils', () => {
       expect(result.name.text).toEqual('main');
     });
   });
+
+  describe('addRoute', () => {
+    let tree: Tree;
+    let context: any;
+
+    beforeEach(() => {
+      context = {
+        warn: jest.fn()
+      };
+      tree = Tree.empty();
+    });
+
+    it('should add links and routes if they are not present', async () => {
+      const sourceCode = `
+import React from 'react';
+const App = () => (
+  <>
+    <h1>Hello</h1>
+  </>
+);
+export default App; 
+      `;
+      tree.create('app.tsx', sourceCode);
+      const source = ts.createSourceFile(
+        'app.tsx',
+        sourceCode,
+        ts.ScriptTarget.Latest,
+        true
+      );
+
+      insert(
+        tree,
+        'app.tsx',
+        utils.addInitialRoutes('app.tsx', source, context)
+      );
+
+      const result = tree.read('app.tsx').toString();
+
+      expect(result).toMatch(/role="navigation"/);
+      expect(result).toMatch(/<Link\s+to="\/page-2"/);
+      expect(result).toMatch(/<Route\s+path="\/page-2"/);
+    });
+
+    it('should update existing routes', async () => {
+      const sourceCode = `
+import React from 'react';
+import { Home } from '@example/home';
+const App = () => (
+  <>
+    <header>
+      <h1>Hello</h1>
+      <div role="navigation">
+        <ul>
+          <li><Link to="/">Home</Link></li>
+        </ul>
+      </div>
+    </header>
+    <p>Hello World!</p>
+    <Route path="/" component={Home}/>
+  </>
+);
+export default App; 
+      `;
+      tree.create('app.tsx', sourceCode);
+      const source = ts.createSourceFile(
+        'app.tsx',
+        sourceCode,
+        ts.ScriptTarget.Latest,
+        true
+      );
+
+      insert(
+        tree,
+        'app.tsx',
+        utils.addRoute(
+          'app.tsx',
+          source,
+          {
+            routePath: '/about',
+            componentName: 'About',
+            moduleName: '@example/about'
+          },
+          context
+        )
+      );
+
+      const result = tree.read('app.tsx').toString();
+
+      expect(result).toMatch(/<li><Link\s+to="\/about"/);
+      expect(result).toMatch(/<Route\s+path="\/about"\s+component={About}/);
+    });
+  });
+
+  describe('addBrowserRouter', () => {
+    let tree: Tree;
+    let context: any;
+
+    beforeEach(() => {
+      context = {
+        warn: jest.fn()
+      };
+      tree = Tree.empty();
+    });
+
+    it('should wrap around App component', () => {
+      const sourceCode = `
+import React from 'react';
+import ReactDOM from 'react-dom';
+import { App } from '@example/my-app';
+ReactDOM.render(<App/>, document.getElementById('root'));
+      `;
+      tree.create('app.tsx', sourceCode);
+      const source = ts.createSourceFile(
+        'app.tsx',
+        sourceCode,
+        ts.ScriptTarget.Latest,
+        true
+      );
+
+      insert(
+        tree,
+        'app.tsx',
+        utils.addBrowserRouter('app.tsx', source, context)
+      );
+
+      const result = tree.read('app.tsx').toString();
+      expect(result).toContain('<BrowserRouter><App/></BrowserRouter>');
+    });
+  });
 });

packages/react/src/utils/ast-utils.ts

@@ -5,17 +5,43 @@ import {
   InsertChange
 } from '@nrwl/workspace/src/utils/ast-utils';
 import * as ts from 'typescript';
+import { noop, SchematicContext } from '@angular-devkit/schematics';
+import { join } from '@angular-devkit/core';
 
-export function addRouter(sourcePath: string, source: ts.SourceFile): Change[] {
-  const jsxClosing = findNodes(source, ts.SyntaxKind.JsxClosingElement);
+export function addInitialRoutes(
+  sourcePath: string,
+  source: ts.SourceFile,
+  context: SchematicContext
+): Change[] {
+  const jsxClosingElements = findNodes(source, [
+    ts.SyntaxKind.JsxClosingElement,
+    ts.SyntaxKind.JsxClosingFragment
+  ]);
+  const outerMostJsxClosing = jsxClosingElements[jsxClosingElements.length - 1];
 
-  const outerMostJsxClosing = jsxClosing[jsxClosing.length - 1];
+  if (!outerMostJsxClosing) {
+    context.logger.warn(
+      `Could not find JSX elements in ${sourcePath}; Skipping insert routes`
+    );
+    return [];
+  }
 
-  const insertRoute = new InsertChange(
+  const insertRoutes = new InsertChange(
     sourcePath,
     outerMostJsxClosing.getStart(),
     `
-    <hr style={{ margin: '36px 0' }}/>
+    {/* START: routes */}
+    {/* These routes and navigation have been generated for you */}
+    {/* Feel free to move and update them to fit your needs */}
+    <br/>
+    <hr/>
+    <br/>
+    <div role="navigation">
+      <ul>
+        <li><Link to="/">Home</Link></li>
+        <li><Link to="/page-2">Page 2</Link></li>
+      </ul>
+    </div>
     <Route
       path="/"
       exact
@@ -30,18 +56,17 @@ export function addRouter(sourcePath: string, source: ts.SourceFile): Change[] {
         <div><Link to="/">Click here to go back to root page.</Link></div>
       )}
     />
+    {/* END: routes */}
     `
   );
 
-  findDefaultExport(source);
-
   return [
     ...addGlobal(
       source,
       sourcePath,
-      `import { Route, Link} from 'react-router-dom';`
+      `import { Route, Link } from 'react-router-dom';`
     ),
-    insertRoute
+    insertRoutes
   ];
 }
 
@@ -49,47 +74,91 @@ export function addRoute(
   sourcePath: string,
   source: ts.SourceFile,
   options: {
-    libName: string;
+    routePath: string;
     componentName: string;
     moduleName: string;
-  }
+  },
+  context: SchematicContext
 ): Change[] {
-  const defaultExport = findDefaultExport(source);
+  const routes = findElements(source, 'Route');
+  const links = findElements(source, 'Link');
 
-  if (!defaultExport) {
-    throw new Error(`Cannot find default export in ${sourcePath}`);
-  }
-
-  const elements = findNodes(
-    defaultExport,
-    ts.SyntaxKind.JsxSelfClosingElement
-  ) as ts.JsxSelfClosingElement[];
-
-  const routes = elements.filter(
-    x =>
-      x.tagName.kind === ts.SyntaxKind.Identifier && x.tagName.text === 'Route'
-  );
-
-  if (routes.length > 0) {
-    const lastRoute = routes[0];
-
-    const addImport = addGlobal(
-      source,
-      sourcePath,
-      `import { ${options.componentName} } from '${options.moduleName}';`
+  if (routes.length === 0) {
+    context.logger.warn(
+      `Could not find <Route/> components in ${sourcePath}; Skipping add route`
     );
-
-    const insertRoute = new InsertChange(
-      sourcePath,
-      lastRoute.getEnd(),
-      `<Route path="/${options.libName}" component={${
-        options.componentName
-      }} />`
-    );
-
-    return [...addImport, insertRoute];
+    return [];
   } else {
-    throw new Error(`Could not find routes in ${sourcePath}`);
+    const changes: Change[] = [];
+    const firstRoute = routes[0];
+    const firstLink = links[0];
+
+    changes.push(
+      ...addGlobal(
+        source,
+        sourcePath,
+        `import { ${options.componentName} } from '${options.moduleName}';`
+      )
+    );
+
+    changes.push(
+      new InsertChange(
+        sourcePath,
+        firstRoute.getEnd(),
+        `<Route path="${options.routePath}" component={${
+          options.componentName
+        }} />`
+      )
+    );
+
+    if (firstLink) {
+      const parentLi = findClosestOpening('li', firstLink);
+      if (parentLi) {
+        changes.push(
+          new InsertChange(
+            sourcePath,
+            parentLi.getEnd(),
+            `<li><Link to="${options.routePath}">${
+              options.componentName
+            }</Link></li>`
+          )
+        );
+      } else {
+        changes.push(
+          new InsertChange(
+            sourcePath,
+            firstLink.parent.getEnd(),
+            `<Link to="${options.routePath}">${options.componentName}</Link>`
+          )
+        );
+      }
+    }
+
+    return changes;
+  }
+}
+
+export function addBrowserRouter(
+  sourcePath: string,
+  source: ts.SourceFile,
+  context: SchematicContext
+): Change[] {
+  const app = findElements(source, 'App')[0];
+  if (app) {
+    return [
+      ...addGlobal(
+        source,
+        sourcePath,
+        `import { BrowserRouter } from 'react-router-dom';`
+      ),
+      new InsertChange(sourcePath, app.getStart(), `<BrowserRouter>`),
+      new InsertChange(sourcePath, app.getEnd(), `</BrowserRouter>`)
+    ];
+  } else {
+    context.logger.warn(
+      `Could not find App component in ${sourcePath}; Skipping add <BrowserRouter>`
+    );
+    return [];
   }
 }
 
@@ -155,6 +224,29 @@ export function findDefaultClassOrFunction(
   );
 }
 
+export function findComponentImportPath(name: string, source: ts.SourceFile) {
+  const allImports = findNodes(
+    source,
+    ts.SyntaxKind.ImportDeclaration
+  ) as ts.ImportDeclaration[];
+  const matching = allImports.filter((i: ts.ImportDeclaration) => {
+    return (
+      i.importClause &&
+      i.importClause.name &&
+      i.importClause.name.getText() === name
+    );
+  });
+
+  if (matching.length === 0) {
+    return null;
+  }
+
+  const appImport = matching[0];
+  return appImport.moduleSpecifier.getText().replace(/['"]/g, '');
+}
+
+// -----------------------------------------------------------------------------
+
 function hasDefaultExportModifier(
   x: ts.ClassDeclaration | ts.FunctionDeclaration
 ) {
@@ -163,3 +255,41 @@ function hasDefaultExportModifier(
     x.modifiers.some(m => m.kind === ts.SyntaxKind.DefaultKeyword)
   );
 }
+
+function findElements(source: ts.SourceFile, tagName: string) {
+  const nodes = findNodes(source, [
+    ts.SyntaxKind.JsxSelfClosingElement,
+    ts.SyntaxKind.JsxOpeningElement
+  ]);
+  return nodes.filter(node => isTag(tagName, node));
+}
+
+function findClosestOpening(tagName: string, node: ts.Node) {
+  if (!node) {
+    return null;
+  }
+
+  if (isTag(tagName, node)) {
+    return node;
+  } else {
+    return findClosestOpening(tagName, node.parent);
+  }
+}
+
+function isTag(tagName: string, node: ts.Node) {
+  if (ts.isJsxOpeningLikeElement(node)) {
+    return (
+      node.tagName.kind === ts.SyntaxKind.Identifier &&
+      node.tagName.text === tagName
+    );
+  }
+
+  if (ts.isJsxElement(node) && node.openingElement) {
+    return (
+      node.openingElement.tagName.kind === ts.SyntaxKind.Identifier &&
+      node.openingElement.tagName.text === tagName
+    );
+  }
+
+  return false;
+}

packages/workspace/src/utils/ast-utils.ts

@@ -45,7 +45,7 @@ function insertAfterLastOccurrence(
 
 export function findNodes(
   node: ts.Node,
-  kind: ts.SyntaxKind,
+  kind: ts.SyntaxKind | ts.SyntaxKind[],
   max = Infinity
 ): ts.Node[] {
   if (!node || max == 0) {
@@ -53,7 +53,10 @@ export function findNodes(
   }
 
   const arr: ts.Node[] = [];
-  if (node.kind === kind) {
+  const hasMatch = Array.isArray(kind)
+    ? kind.includes(node.kind)
+    : node.kind === kind;
+  if (hasMatch) {
     arr.push(node);
     max--;
   }