implement python support, create some helper classes

Author: florianbgt

.gitignore

@@ -2,4 +2,6 @@ node_modules
 dist
 .env
 binaries
-napi_dist
+napi_dist
+__pycache__
+.ruff_cache

package-lock.json

@@ -82,6 +82,7 @@
     "prompts": "^2.4.2",
     "tree-sitter": "^0.21.1",
     "tree-sitter-javascript": "^0.23.0",
+    "tree-sitter-python": "^0.23.4",
     "tree-sitter-typescript": "^0.23.0",
     "uuid": "^10.0.0",
     "yargs": "^17.7.2",

packages/cli/package.json

@@ -1,14 +1,11 @@
 import { z } from "zod";
-import {
-  getDependencyTree,
-  getEndpontsFromTree,
-} from "../helper/dependencyTree";
+import DependencyTreeManager from "../helper/dependencyTree";
 import { scanSchema } from "./helpers/validation";
 
 export function scan(payload: z.infer<typeof scanSchema>) {
-  const tree = getDependencyTree(payload.entrypointPath);
-
-  const endpoints = getEndpontsFromTree(tree);
-
+  const dependencyTreeManager = new DependencyTreeManager(
+    payload.entrypointPath,
+  );
+  const endpoints = dependencyTreeManager.getEndponts();
   return { endpoints };
 }

packages/cli/src/api/scan.ts

@@ -1,12 +1,8 @@
 import fs from "fs";
 import path from "path";
-import {
-  getDependencyTree,
-  getEndpontsFromTree,
-  getGroupsFromEndpoints,
-} from "../helper/dependencyTree";
+import DependencyTreeManager from "../helper/dependencyTree";
 import { cleanupOutputDir, createOutputDir } from "../helper/file";
-import { createSplit } from "../helper/split";
+import SplitRunner from "../helper/split";
 import { splitSchema } from "./helpers/validation";
 import { z } from "zod";
 import { GroupMap } from "../helper/types";
@@ -16,25 +12,38 @@ export function split(payload: z.infer<typeof splitSchema>) {
   const groupMap: GroupMap = {};
 
   // Get the dependency tree
-  const tree = getDependencyTree(payload.entrypointPath);
+  const dependencyTreeManager = new DependencyTreeManager(
+    payload.entrypointPath,
+  );
 
   const outputDir = payload.outputDir || path.dirname(payload.entrypointPath);
 
   // Clean up and prepare the output directory
   cleanupOutputDir(outputDir);
   createOutputDir(outputDir);
 
-  // Iterate over the tree and get endpoints
-  const endpoints = getEndpontsFromTree(tree);
-
-  // Get groups from the endpoints
-  const groups = getGroupsFromEndpoints(endpoints);
+  // Get groups from the dependency tree
+  const groups = dependencyTreeManager.getGroups();
 
   // Process each group for splitting
   groups.forEach((group, index) => {
-    // Clone the tree to avoid mutation of the original tree
-    const treeClone = structuredClone(tree);
-    createSplit(treeClone, group, outputDir, payload.entrypointPath, index);
+    console.log(JSON.stringify(dependencyTreeManager.dependencyTree, null, 2));
+
+    const splitRunner = new SplitRunner(dependencyTreeManager, group);
+    const files = splitRunner.run();
+
+    const targetDir = path.dirname(payload.entrypointPath);
+    const annotationDirectory = path.join(outputDir, index.toString());
+
+    files.forEach((file) => {
+      const relativeFileNamePath = path.relative(targetDir, file.path);
+      const destinationPath = path.join(
+        annotationDirectory,
+        relativeFileNamePath,
+      );
+      fs.mkdirSync(path.dirname(destinationPath), { recursive: true });
+      fs.writeFileSync(destinationPath, file.sourceCode, "utf8");
+    });
   });
 
   // Store the processed annotations in the output directory

packages/cli/src/api/split.ts

@@ -5,20 +5,19 @@ import {
   parseNanoApiAnnotation,
   updateCommentFromAnnotation,
 } from "../helper/annotations";
-import {
-  getDependencyTree,
-  getEndpontsFromTree,
-} from "../helper/dependencyTree";
+import DependencyTreeManager from "../helper/dependencyTree";
 import Parser from "tree-sitter";
 import { getParserLanguageFromFile } from "../helper/treeSitter";
 import { replaceIndexesFromSourceCode } from "../helper/cleanup";
 import { getJavascriptAnnotationNodes } from "../helper/languages/javascript/annotations";
 import { getTypescriptAnnotationNodes } from "../helper/languages/typescript/annotations";
 
 export function sync(payload: z.infer<typeof syncSchema>) {
-  const tree = getDependencyTree(payload.entrypointPath);
+  const dependencyTreeManager = new DependencyTreeManager(
+    payload.entrypointPath,
+  );
 
-  const endpoints = getEndpontsFromTree(tree);
+  const endpoints = dependencyTreeManager.getEndponts();
 
   const updatedEndpoints = endpoints.map((endpoint) => {
     const updatedEndpoint = payload.endpoints.find(

packages/cli/src/api/sync.ts

@@ -1,5 +1,5 @@
 import fs from "fs";
-import { getDependencyTree } from "../helper/dependencyTree";
+import DependencyTreeManager from "../helper/dependencyTree";
 import OpenAI from "openai";
 import { DependencyTree } from "../helper/types";
 
@@ -8,7 +8,7 @@ export default async function annotateOpenAICommandHandler(
   openAIApiKey: string, // OpenAI API key
 ) {
   console.info("Annotating program...");
-  const tree = getDependencyTree(entrypoint);
+  const dependencyTreeManager = new DependencyTreeManager(entrypoint);
 
   const openAIConfig = {
     apiKey: openAIApiKey,
@@ -24,7 +24,7 @@ export default async function annotateOpenAICommandHandler(
   messages.push({
     role: "user",
     content: `The entrypoint of my application is: ${entrypoint}.
-    Here is the dependency tree of the entrypoint of: ${JSON.stringify(tree)}`,
+    Here is the dependency tree of the entrypoint of: ${JSON.stringify(dependencyTreeManager.dependencyTree)}`,
   });
 
   messages.push({
@@ -50,7 +50,7 @@ export default async function annotateOpenAICommandHandler(
       iterateOverTree(child);
     });
   }
-  iterateOverTree(tree);
+  iterateOverTree(dependencyTreeManager.dependencyTree);
 
   messages.push({
     role: "user",

packages/cli/src/commands/annotate.ts

@@ -1,12 +1,8 @@
 import path from "path";
 import fs from "fs";
-import {
-  getDependencyTree,
-  getEndpontsFromTree,
-  getGroupsFromEndpoints,
-} from "../helper/dependencyTree";
+import DependencyTreeManager from "../helper/dependencyTree";
 import { cleanupOutputDir, createOutputDir } from "../helper/file";
-import { createSplit } from "../helper/split";
+import SplitRunner from "../helper/split";
 import { GroupMap } from "../helper/types";
 
 export default function splitCommandHandler(
@@ -15,22 +11,31 @@ export default function splitCommandHandler(
 ) {
   const groupMap: GroupMap = {};
 
-  const tree = getDependencyTree(entrypointPath);
+  const dependencyTreeManager = new DependencyTreeManager(entrypointPath);
 
   cleanupOutputDir(outputDir);
   createOutputDir(outputDir);
 
-  // Iterate over the tree and get endpoints
-  const endpoints = getEndpontsFromTree(tree);
-
-  // Get groups from the endpoints
-  const groups = getGroupsFromEndpoints(endpoints);
+  // Get groups from the dependency tree
+  const groups = dependencyTreeManager.getGroups();
 
   // Process each group for splitting
   groups.forEach((group, index) => {
-    // Clone the tree to avoid mutation of the original tree
-    const treeClone = structuredClone(tree);
-    createSplit(treeClone, group, outputDir, entrypointPath, index);
+    const splitRunner = new SplitRunner(dependencyTreeManager, group);
+    const files = splitRunner.run();
+
+    const targetDir = path.dirname(entrypointPath);
+    const annotationDirectory = path.join(outputDir, index.toString());
+
+    files.forEach((file) => {
+      const relativeFileNamePath = path.relative(targetDir, file.path);
+      const destinationPath = path.join(
+        annotationDirectory,
+        relativeFileNamePath,
+      );
+      fs.mkdirSync(path.dirname(destinationPath), { recursive: true });
+      fs.writeFileSync(destinationPath, file.sourceCode, "utf8");
+    });
   });
 
   // Store the processed annotations in the output directory