Complete function rename; refactored notebook logic out of main function.

brdskggs · brdskggs · commit db49f5dec6cb · 2025-05-07T13:01:39.000-04:00
diff --git a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts
@@ -213,15 +213,15 @@ describe('editorContext', function () {
         })
     })
 
-    describe('extractCellsSliceContext', function () {
+    describe('getNotebookCellsSliceContext', function () {
         it('Should extract content from cells in reverse order up to maxLength from prefix cells', function () {
             const mockCells = [
                 createNotebookCell(createMockDocument('First cell content')),
                 createNotebookCell(createMockDocument('Second cell content')),
                 createNotebookCell(createMockDocument('Third cell content')),
             ]
 
-            const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', false)
+            const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', false)
             assert.strictEqual(result, 'First cell content\nSecond cell content\nThird cell content\n')
         })
 
@@ -232,7 +232,7 @@ describe('editorContext', function () {
                 createNotebookCell(createMockDocument('Third cell content')),
             ]
 
-            const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', true)
+            const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', true)
             assert.strictEqual(result, 'First cell content\nSecond cell content\nThird cell content\n')
         })
 
@@ -244,7 +244,7 @@ describe('editorContext', function () {
                 createNotebookCell(createMockDocument('Fourth')),
             ]
             // Should only include part of second cell and the last two cells
-            const result = EditorContext.extractCellsSliceContext(mockCells, 15, 'python', false)
+            const result = EditorContext.getNotebookCellsSliceContext(mockCells, 15, 'python', false)
             assert.strictEqual(result, 'd\nThird\nFourth\n')
         })
 
@@ -257,17 +257,17 @@ describe('editorContext', function () {
             ]
 
             // Should only include first cell and part of second cell
-            const result = EditorContext.extractCellsSliceContext(mockCells, 15, 'python', true)
+            const result = EditorContext.getNotebookCellsSliceContext(mockCells, 15, 'python', true)
             assert.strictEqual(result, 'First\nSecond\nTh')
         })
 
         it('Should handle empty cells array from prefix cells', function () {
-            const result = EditorContext.extractCellsSliceContext([], 100, 'python', false)
+            const result = EditorContext.getNotebookCellsSliceContext([], 100, 'python', false)
             assert.strictEqual(result, '')
         })
 
         it('Should handle empty cells array from suffix cells', function () {
-            const result = EditorContext.extractCellsSliceContext([], 100, 'python', true)
+            const result = EditorContext.getNotebookCellsSliceContext([], 100, 'python', true)
             assert.strictEqual(result, '')
         })
 
@@ -276,7 +276,7 @@ describe('editorContext', function () {
                 createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup),
                 createNotebookCell(createMockDocument('def example():\n    return "test"')),
             ]
-            const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', false)
+            const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', false)
             assert.strictEqual(result, '# # Heading\n# This is markdown\ndef example():\n    return "test"\n')
         })
 
@@ -286,7 +286,7 @@ describe('editorContext', function () {
                 createNotebookCell(createMockDocument('def example():\n    return "test"')),
             ]
 
-            const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', true)
+            const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', true)
             assert.strictEqual(result, '# # Heading\n# This is markdown\ndef example():\n    return "test"\n')
         })
 
@@ -295,7 +295,7 @@ describe('editorContext', function () {
                 createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup),
                 createNotebookCell(createMockDocument('def example():\n    return "test"')),
             ]
-            const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'java', false)
+            const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'java', false)
             assert.strictEqual(result, '// # Heading\n// This is markdown\n// def example():\n//     return "test"\n')
         })
 
@@ -305,7 +305,7 @@ describe('editorContext', function () {
                 createNotebookCell(createMockDocument('println(1 + 1);', 'somefile.ipynb', 'java')),
             ]
 
-            const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'java', true)
+            const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'java', true)
             assert.strictEqual(result, '// # Heading\n// This is markdown\nprintln(1 + 1);\n')
         })
 
@@ -317,7 +317,7 @@ describe('editorContext', function () {
                 ),
                 createNotebookCell(createMockDocument('def example():\n    return "test"')),
             ]
-            const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', false)
+            const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', false)
             assert.strictEqual(result, '# println(1 + 1);\ndef example():\n    return "test"\n')
         })
 
@@ -329,7 +329,7 @@ describe('editorContext', function () {
                 ),
                 createNotebookCell(createMockDocument('def example():\n    return "test"')),
             ]
-            const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', true)
+            const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', true)
             assert.strictEqual(result, '# println(1 + 1);\ndef example():\n    return "test"\n')
         })
     })
diff --git a/packages/core/src/codewhisperer/util/editorContext.ts b/packages/core/src/codewhisperer/util/editorContext.ts
@@ -29,7 +29,52 @@ const languageCommentChars: Record<string, string> = {
     java: '// ',
 }
 
-export function getNotebookCellContext(cell: vscode.NotebookCell, referenceLanguage?: string): string {
+function getEnclosingNotebook(editor: vscode.TextEditor): vscode.NotebookDocument | undefined {
+    // For notebook cells, find the existing notebook with a cell that matches the current editor.
+    return vscode.workspace.notebookDocuments.find(
+        (nb) =>
+            nb.notebookType === 'jupyter-notebook' && nb.getCells().some((cell) => cell.document === editor.document)
+    )
+}
+
+export function extractNotebookContext(
+    notebook: vscode.NotebookDocument,
+    editor: vscode.TextEditor,
+    languageName: string,
+    caretLeftFileContext: string,
+    caretRightFileContext: string
+) {
+    // Expand the context for a cell inside of a noteboo with whatever text fits from the preceding and subsequent cells
+    const allCells = notebook.getCells()
+    const cellIndex = allCells.findIndex((cell) => cell.document === editor.document)
+    // Extract text from prior cells if there is enough room in left file context
+    if (caretLeftFileContext.length < CodeWhispererConstants.charactersLimit - 1) {
+        const leftCellsText = getNotebookCellsSliceContext(
+            allCells.slice(0, cellIndex),
+            CodeWhispererConstants.charactersLimit - (caretLeftFileContext.length + 1),
+            languageName,
+            true
+        )
+        if (leftCellsText.length > 0) {
+            caretLeftFileContext = addNewlineIfMissing(leftCellsText) + caretLeftFileContext
+        }
+    }
+    // Extract text from subsequent cells if there is enough room in right file context
+    if (caretRightFileContext.length < CodeWhispererConstants.charactersLimit - 1) {
+        const rightCellsText = getNotebookCellsSliceContext(
+            allCells.slice(cellIndex + 1),
+            CodeWhispererConstants.charactersLimit - (caretRightFileContext.length + 1),
+            languageName,
+            false
+        )
+        if (rightCellsText.length > 0) {
+            caretRightFileContext = addNewlineIfMissing(caretRightFileContext) + rightCellsText
+        }
+    }
+    return { caretLeftFileContext, caretRightFileContext }
+}
+
+export function extractSingleCellContext(cell: vscode.NotebookCell, referenceLanguage?: string): string {
     // Extract the text verbatim if the cell is code and the cell has the same language.
     // Otherwise, add the correct comment string for the refeference language
     const cellText = cell.document.getText()
@@ -57,6 +102,8 @@ export function getNotebookCellsSliceContext(
     referenceLanguage: string,
     fromStart: boolean
 ): string {
+    // Extract context from array of notebook cells that fits inside `maxLength` characters,
+    // from either the start or the end of the array.
     let output: string[] = []
     if (!fromStart) {
         cells = cells.reverse()
@@ -112,40 +159,15 @@ export function extractContextForCodeWhisperer(editor: vscode.TextEditor): codew
             runtimeLanguageContext.normalizeLanguage(editor.document.languageId) ?? editor.document.languageId
     }
     if (editor.document.uri.scheme === 'vscode-notebook-cell') {
-        // For notebook cells, first find the existing notebook with a cell that matches the current editor.
-        const notebook = vscode.workspace.notebookDocuments.find(
-            (nb) =>
-                nb.notebookType === 'jupyter-notebook' &&
-                nb.getCells().some((cell) => cell.document === editor.document)
-        )
+        const notebook = getEnclosingNotebook(editor)
         if (notebook) {
-            const allCells = notebook.getCells()
-            const cellIndex = allCells.findIndex((cell) => cell.document === editor.document)
-
-            // Extract text from prior cells if there is enough room in left file context
-            if (caretLeftFileContext.length < CodeWhispererConstants.charactersLimit - 1) {
-                const leftCellsText = extractCellsSliceContext(
-                    allCells.slice(0, cellIndex),
-                    CodeWhispererConstants.charactersLimit - (caretLeftFileContext.length + 1),
-                    languageName,
-                    true
-                )
-                if (leftCellsText.length > 0) {
-                    caretLeftFileContext = addNewlineIfMissing(leftCellsText) + caretLeftFileContext
-                }
-            }
-            // Extract text from subsequent cells if there is enough room in right file context
-            if (caretRightFileContext.length < CodeWhispererConstants.charactersLimit - 1) {
-                const rightCellsText = extractCellsSliceContext(
-                    allCells.slice(cellIndex + 1),
-                    CodeWhispererConstants.charactersLimit - (caretRightFileContext.length + 1),
-                    languageName,
-                    false
-                )
-                if (rightCellsText.length > 0) {
-                    caretRightFileContext = addNewlineIfMissing(caretRightFileContext) + rightCellsText
-                }
-            }
+            ;({ caretLeftFileContext, caretRightFileContext } = extractNotebookContext(
+                notebook,
+                editor,
+                languageName,
+                caretLeftFileContext,
+                caretRightFileContext
+            ))
         }
     }
 
