feat: basic js ast explorer (#155)

* stash in commit, because i forget about stashes

* feat: basic js ast explorer

* add variable_declarator condition to getCalls

* refactor: format

* fix(deps): add missing dependencies

---------

Co-authored-by: Oliver Eyton-Williams <[email protected]>

Author: ShaunSHamilton

lib/__tests__/tower.test.ts

@@ -0,0 +1,94 @@
+import { Tower, generate } from "../index";
+
+const code = `
+import { readFile } from "fs/promises";
+
+async function main() {
+  doSyncStuff();
+
+  const a = await doAsyncStuff(1);
+  console.log(a);
+}
+
+// Test
+async function doAsyncStuff(num) {
+  // Another test
+  const file = await readFile("file_loc" + num, "utf-8");
+  return file;
+}
+
+const a = [];
+
+const b = {
+  c: () => {
+    const c = 1;
+  },
+  "example-a": {
+    a: 1
+  }
+}
+
+const { c, d, e } = b;
+
+const [f, g] = [c, d];
+
+let h, i;
+
+a.map(function (e) {});
+
+function doSyncStuff() {
+  console.log("stuff");
+}
+
+main();
+`;
+
+const t = new Tower(code);
+
+describe("tower", () => {
+  describe("getFunction", () => {
+    it("works", () => {
+      const main = t.getFunction("main").generate;
+      expect(main).toEqual(
+        `async function main() {
+  doSyncStuff();
+  const a = await doAsyncStuff(1);
+  console.log(a);
+}
+
+// Test`,
+      );
+    });
+  });
+  describe("getVariable", () => {
+    it("works", () => {
+      const a = t.getFunction("main").getVariable("a").generate;
+      expect(a).toEqual("const a = await doAsyncStuff(1);");
+      const b = t.getVariable("b").generate;
+      expect(b).toEqual(`const b = {
+  c: () => {
+    const c = 1;
+  },
+  "example-a": {
+    a: 1
+  }
+};`);
+      const file = t.getFunction("doAsyncStuff").getVariable("file").generate;
+      expect(file).toEqual(
+        '// Another test\nconst file = await readFile("file_loc" + num, "utf-8");',
+      );
+    });
+  });
+  describe("getCalls", () => {
+    it("works", () => {
+      const aMap = t.getCalls("a.map");
+      expect(aMap).toHaveLength(1);
+      const map = aMap.at(0);
+      // @ts-expect-error - expression does exist.
+      const argumes = map?.ast.expression?.arguments;
+      expect(generate(argumes.at(0), { compact: true }).code).toEqual(
+        "function(e){}",
+      );
+    });
+  });
+});

lib/class/tower.ts

@@ -0,0 +1,142 @@
+import { parse, ParserOptions } from "@babel/parser";
+import generate from "@babel/generator";
+import {
+  ExpressionStatement,
+  FunctionDeclaration,
+  is,
+  Node,
+  VariableDeclaration,
+} from "@babel/types";
+
+export { generate };
+
+export class Tower<T extends Node> {
+  public ast: Node;
+  constructor(stringOrAST: string | T, options?: Partial<ParserOptions>) {
+    if (typeof stringOrAST === "string") {
+      const parsedThing = parse(stringOrAST, {
+        sourceType: "module",
+        ...options,
+      });
+      this.ast = parsedThing.program;
+    } else {
+      this.ast = stringOrAST;
+    }
+  }
+
+  // Get all the given types at the current scope
+  private getType<T extends Node>(type: string, name: string): Tower<T> {
+    const body = this.extractBody(this.ast);
+    const ast = body.find((node) => {
+      if (node.type === type) {
+        if (is("FunctionDeclaration", node)) {
+          return node.id?.name === name;
+        }
+
+        if (is("VariableDeclaration", node)) {
+          const variableDeclarator = node.declarations[0];
+          if (!is("VariableDeclarator", variableDeclarator)) {
+            return false;
+          }
+
+          const identifier = variableDeclarator.id;
+          if (!is("Identifier", identifier)) {
+            return false;
+          }
+
+          return identifier.name === name;
+        }
+      }
+
+      return false;
+    });
+    if (!ast) {
+      throw new Error(`No AST found with name ${name}`);
+    }
+
+    assertIsType<T>(ast);
+    return new Tower<T>(ast);
+  }
+
+  public getFunction(name: string): Tower<FunctionDeclaration> {
+    return this.getType("FunctionDeclaration", name);
+  }
+
+  public getVariable(name: string): Tower<VariableDeclaration> {
+    return this.getType("VariableDeclaration", name);
+  }
+
+  public getCalls(callSite: string): Array<Tower<ExpressionStatement>> {
+    const body = this.extractBody(this.ast);
+    const calls = body.filter((node) => {
+      if (is("ExpressionStatement", node)) {
+        const expression = node.expression;
+        if (is("CallExpression", expression)) {
+          const callee = expression.callee;
+
+          switch (callee.type) {
+            case "Identifier":
+              return callee.name === callSite;
+            case "MemberExpression":
+              return generate(callee).code === callSite;
+            default:
+              return true;
+          }
+        }
+      }
+
+      if (is("VariableDeclarator", node)) {
+        const init = node.init;
+        if (is("CallExpression", init)) {
+          const callee = init.callee;
+
+          switch (callee.type) {
+            case "Identifier":
+              return callee.name === callSite;
+            case "MemberExpression":
+              return generate(callee).code === callSite;
+            default:
+              return true;
+          }
+        }
+      }
+
+      return false;
+    });
+    assertIsType<ExpressionStatement[]>(calls);
+    return calls.map((call) => new Tower<ExpressionStatement>(call));
+  }
+
+  private extractBody(ast: Node): Node[] {
+    switch (ast.type) {
+      case "Program":
+        return ast.body;
+      case "FunctionDeclaration":
+        return ast.body.body;
+      case "VariableDeclaration":
+        return ast.declarations;
+      case "ArrowFunctionExpression":
+        // eslint-disable-next-line no-case-declarations
+        const blockStatement = ast.body;
+        if (is("BlockStatement", blockStatement)) {
+          return blockStatement.body;
+        }
+
+        throw new Error(`Unimplemented for ${ast.type}`);
+      default:
+        throw new Error(`Unimplemented for ${ast.type}`);
+    }
+  }
+
+  public get generate(): string {
+    return generate(this.ast).code;
+  }
+
+  public get compact(): string {
+    return generate(this.ast, { compact: true }).code;
+  }
+}
+
+function assertIsType<T extends Node | Node[]>(
+  ast: Node | Node[],
+): asserts ast is T {}

lib/index.ts

@@ -1,5 +1,6 @@
 import { strip } from "./strip";
 import astHelpers from "../python/py_helpers.py";
+export { Tower, generate } from "./class/tower";
 export { Babeliser } from "./class/babeliser";
 
 declare global {

package.json

@@ -76,15 +76,17 @@
     "build": "NODE_ENV=production webpack",
     "clean": "rm -rf dist",
     "lint": "eslint lib --max-warnings 0 && prettier lib --check && tsc --noEmit",
+    "format": "prettier --write lib",
     "test": "jest",
     "prepublishOnly": "pnpm clean && pnpm build",
     "webpack": "webpack"
   },
   "repository": "[email protected]:freeCodeCamp/curriculum-helpers.git",
   "license": "BSD-3-Clause",
   "dependencies": {
-    "@babel/generator": "7.x",
-    "@babel/parser": "7.x",
+    "@babel/generator": "7.26.9",
+    "@babel/parser": "7.26.9",
+    "@babel/types": "7.26.9",
     "browserify": "^17.0.0"
   }
 }