From 58e211d09f6b720a3a40d42b773a775ea3661242 Mon Sep 17 00:00:00 2001
From: Jacob Lee <jacoblee93@gmail.com>
Date: Tue, 30 Jul 2024 15:20:38 -0700
Subject: [PATCH] ci[patch]: Allow duplicate imports in notebooks (#6283)

* Allow duplicate imports in notebooks

* Add dev dep

* Format
---
 docs/core_docs/package.json                 |  1 +
 docs/core_docs/scripts/validate_notebook.ts | 45 ++++++++++++++++-----
 yarn.lock                                   | 32 ++++++++++++++-
 3 files changed, 68 insertions(+), 10 deletions(-)

diff --git a/docs/core_docs/package.json b/docs/core_docs/package.json
index e1f6de183330..ec24faff8686 100644
--- a/docs/core_docs/package.json
+++ b/docs/core_docs/package.json
@@ -71,6 +71,7 @@
     "rimraf": "^5.0.1",
     "supabase": "^1.148.6",
     "swc-loader": "^0.2.3",
+    "ts-morph": "^23.0.0",
     "tsx": "^3.12.3",
     "typedoc": "^0.24.4",
     "typedoc-plugin-markdown": "next",
diff --git a/docs/core_docs/scripts/validate_notebook.ts b/docs/core_docs/scripts/validate_notebook.ts
index 00ba61aa30d0..6f26d940be3a 100644
--- a/docs/core_docs/scripts/validate_notebook.ts
+++ b/docs/core_docs/scripts/validate_notebook.ts
@@ -1,17 +1,44 @@
 import * as fs from "node:fs";
 import * as ts from "typescript";
+import { Project } from "ts-morph";
 
 export function extract(filepath: string) {
   const cells = JSON.parse(fs.readFileSync(filepath).toString()).cells;
-  const code = cells
-    .map((cell: Record<string, any>) => {
-      if (cell.cell_type === "code") {
-        return cell.source.join("");
-      }
-      return "";
-    })
-    .join("\n");
-  return code;
+  const project = new Project({ useInMemoryFileSystem: true });
+  const sourceFile = project.createSourceFile("temp.ts", "");
+
+  cells.forEach((cell: Record<string, any>) => {
+    if (cell.cell_type === "code") {
+      sourceFile.addStatements(cell.source.join(""));
+    }
+  });
+
+  // Deduplicate imports
+  const importDeclarations = sourceFile.getImportDeclarations();
+  const uniqueImports = new Map<string, Set<string>>();
+
+  importDeclarations.forEach((importDecl) => {
+    const moduleSpecifier = importDecl.getModuleSpecifierValue();
+    if (!uniqueImports.has(moduleSpecifier)) {
+      uniqueImports.set(moduleSpecifier, new Set());
+    }
+    importDecl.getNamedImports().forEach((namedImport) => {
+      uniqueImports.get(moduleSpecifier)!.add(namedImport.getText());
+    });
+  });
+
+  // Remove all existing imports
+  importDeclarations.forEach((importDecl) => importDecl.remove());
+
+  // Add deduplicated imports at the top
+  uniqueImports.forEach((namedImports, moduleSpecifier) => {
+    sourceFile.addImportDeclaration({
+      moduleSpecifier,
+      namedImports: Array.from(namedImports),
+    });
+  });
+
+  return sourceFile.getFullText();
 }
 
 let [pathname, ...args] = process.argv.slice(2);
diff --git a/yarn.lock b/yarn.lock
index 03dec4a5a6c7..d1dbf36543be 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -17874,6 +17874,18 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@ts-morph/common@npm:~0.24.0":
+  version: 0.24.0
+  resolution: "@ts-morph/common@npm:0.24.0"
+  dependencies:
+    fast-glob: ^3.3.2
+    minimatch: ^9.0.4
+    mkdirp: ^3.0.1
+    path-browserify: ^1.0.1
+  checksum: 793bc8a47c93ab55c6c036f94480d3b0e948661aef4bb7dbc29279b1dda2fc4fce809a88e221537867a313541842e12d1ecbd32b4769688abe1303807ec09db6
+  languageName: node
+  linkType: hard
+
 "@tsconfig/recommended@npm:^1.0.2":
   version: 1.0.2
   resolution: "@tsconfig/recommended@npm:1.0.2"
@@ -22322,6 +22334,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"code-block-writer@npm:^13.0.1":
+  version: 13.0.2
+  resolution: "code-block-writer@npm:13.0.2"
+  checksum: 8052ae6f27ef73366bd5df04b6f9beced493261fcaef5cbd0f3853644b5e0aa5af18d099b96448be88ea3d000c7b180207d371044edd9fcf98fea22c9f8ba3a1
+  languageName: node
+  linkType: hard
+
 "cohere-ai@npm:>=6.0.0":
   version: 6.2.2
   resolution: "cohere-ai@npm:6.2.2"
@@ -22867,6 +22886,7 @@ __metadata:
     rimraf: ^5.0.1
     supabase: ^1.148.6
     swc-loader: ^0.2.3
+    ts-morph: ^23.0.0
     tsx: ^3.12.3
     typedoc: ^0.24.4
     typedoc-plugin-markdown: next
@@ -32994,7 +33014,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"minimatch@npm:^9.0.5":
+"minimatch@npm:^9.0.4, minimatch@npm:^9.0.5":
   version: 9.0.5
   resolution: "minimatch@npm:9.0.5"
   dependencies:
@@ -39914,6 +39934,16 @@ __metadata:
   languageName: node
   linkType: hard
 
+"ts-morph@npm:^23.0.0":
+  version: 23.0.0
+  resolution: "ts-morph@npm:23.0.0"
+  dependencies:
+    "@ts-morph/common": ~0.24.0
+    code-block-writer: ^13.0.1
+  checksum: 3282eb0f8bd4577770874736c3259b97501da9a86137160b5d68f106b7848ea7b1fbccf9e198a3d930ec40c993e9951d4bfae31e2562dac8f3de0d7bb0e23615
+  languageName: node
+  linkType: hard
+
 "ts-type@npm:>=2":
   version: 3.0.1
   resolution: "ts-type@npm:3.0.1"