From 96efe70063a20023a86ae6cd4806c7bfb229a5aa Mon Sep 17 00:00:00 2001
From: jacoblee93 <jacoblee93@gmail.com>
Date: Thu, 1 Aug 2024 10:34:36 -0700
Subject: [PATCH] Fix notebook lint for namespace imports

---
 docs/core_docs/scripts/validate_notebook.ts | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/docs/core_docs/scripts/validate_notebook.ts b/docs/core_docs/scripts/validate_notebook.ts
index 776684ebc73f..7afdd366cc4a 100644
--- a/docs/core_docs/scripts/validate_notebook.ts
+++ b/docs/core_docs/scripts/validate_notebook.ts
@@ -17,7 +17,7 @@ export function extract(filepath: string) {
   const importDeclarations = sourceFile.getImportDeclarations();
   const uniqueImports = new Map<
     string,
-    { default?: string; named: Set<string> }
+    { default?: string; namespace?: string; named: Set<string> }
   >();
 
   importDeclarations.forEach((importDecl) => {
@@ -29,6 +29,10 @@ export function extract(filepath: string) {
     if (defaultImport) {
       uniqueImports.get(moduleSpecifier)!.default = defaultImport.getText();
     }
+    const namespaceImport = importDecl.getNamespaceImport();
+    if (namespaceImport) {
+      uniqueImports.get(moduleSpecifier)!.namespace = namespaceImport.getText();
+    }
     importDecl.getNamedImports().forEach((namedImport) => {
       uniqueImports.get(moduleSpecifier)!.named.add(namedImport.getText());
     });
@@ -39,10 +43,11 @@ export function extract(filepath: string) {
 
   // Add deduplicated imports at the top
   uniqueImports.forEach(
-    ({ default: defaultImport, named }, moduleSpecifier) => {
+    ({ default: defaultImport, namespace, named }, moduleSpecifier) => {
       sourceFile.addImportDeclaration({
         moduleSpecifier,
         defaultImport,
+        namespaceImport: namespace,
         namedImports: Array.from(named),
       });
     }