diff --git a/packages/jupyter-ai/src/contexts/selection-context.tsx b/packages/jupyter-ai/src/contexts/selection-context.tsx index 3c7a78fca..7101e537e 100644 --- a/packages/jupyter-ai/src/contexts/selection-context.tsx +++ b/packages/jupyter-ai/src/contexts/selection-context.tsx @@ -1,5 +1,6 @@ import React, { useCallback, useContext, useEffect, useState } from 'react'; -import { Selection, SelectionWatcher } from '../selection-watcher'; +import type { Selection } from '../selection-utils'; +import { SelectionWatcher } from '../selection-watcher'; const SelectionContext = React.createContext< [Selection | null, (value: Selection) => unknown] @@ -36,7 +37,9 @@ export function SelectionContextProvider({ const replaceSelection = useCallback( (value: Selection) => { - selectionWatcher.replaceSelection(value); + if (value.type === 'text') { + selectionWatcher.replaceSelection(value); + } }, [selectionWatcher] ); diff --git a/packages/jupyter-ai/src/selection-utils.ts b/packages/jupyter-ai/src/selection-utils.ts new file mode 100644 index 000000000..bdaca5041 --- /dev/null +++ b/packages/jupyter-ai/src/selection-utils.ts @@ -0,0 +1,172 @@ +/** + * Contains various utility functions used by the SelectionWatcher class. + */ +import { DocumentWidget } from '@jupyterlab/docregistry'; +import { Notebook } from '@jupyterlab/notebook'; +import { FileEditor } from '@jupyterlab/fileeditor'; +import { CodeEditor } from '@jupyterlab/codeeditor'; +import { CodeMirrorEditor } from '@jupyterlab/codemirror'; +import { Widget } from '@lumino/widgets'; + +export type BaseSelection = { + /** + * Type of selection. + */ + type: T; + /** + * The ID of the document widget in which the selection was made. + */ + widgetId: string; + /** + * The ID of the cell in which the selection was made, if the original widget + * was a notebook. + */ + cellId?: string; +}; + +export type TextSelection = BaseSelection<'text'> & + CodeEditor.ITextSelection & { + /** + * The text within the selection as a string. + */ + text: string; + }; + +/** + * Represents a selection of one or more cells in a notebook. + */ +export type CellSelection = BaseSelection<'cell'> & { + /** + * The text within the selected cells represented as a string. Multiple cells + * are joined with newlines. + */ + text: string; +}; + +export type Selection = TextSelection | CellSelection; + +/** + * Gets the index of the cell associated with `cellId`. + */ +export function getCellIndex(notebook: Notebook, cellId: string): number { + const idx = notebook.model?.sharedModel.cells.findIndex( + cell => cell.getId() === cellId + ); + return idx === undefined ? -1 : idx; +} + +/** + * Gets the editor instance used by a document widget. Returns `null` if unable. + */ +export function getEditor(widget: Widget | null): CodeMirrorEditor | null { + if (!(widget instanceof DocumentWidget)) { + return null; + } + + let editor: CodeEditor.IEditor | undefined; + const { content } = widget; + + if (content instanceof FileEditor) { + editor = content.editor; + } else if (content instanceof Notebook) { + editor = content.activeCell?.editor; + } + + if (!(editor instanceof CodeMirrorEditor)) { + return null; + } + + return editor; +} + +/** + * Gets a TextSelection object from a document widget. Returns `null` if unable. + */ +function getTextSelection(widget: DocumentWidget): Selection | null { + const editor = getEditor(widget); + // widget type check is redundant but hints the type to TypeScript + if (!editor) { + return null; + } + + let cellId: string | undefined = undefined; + if (widget.content instanceof Notebook) { + cellId = widget.content.activeCell?.model.id; + } + + let { start, end, ...selectionObj } = editor.getSelection(); + const startOffset = editor.getOffsetAt(start); + const endOffset = editor.getOffsetAt(end); + const text = editor.model.sharedModel + .getSource() + .substring(startOffset, endOffset); + + // ensure start <= end + // required for editor.model.sharedModel.updateSource() + if (startOffset > endOffset) { + [start, end] = [end, start]; + } + + return { + type: 'text', + ...selectionObj, + start, + end, + text, + widgetId: widget.id, + ...(cellId && { + cellId + }) + }; +} + +function getSelectedCells(notebook: Notebook) { + return notebook.widgets.filter(cell => notebook.isSelected(cell)); +} + +function getCellSelection(widget: DocumentWidget): CellSelection | null { + if (!(widget.content instanceof Notebook)) { + return null; + } + + const notebook = widget.content; + const selectedCells = getSelectedCells(notebook); + console.log({ selectedCells }); + + const text = selectedCells.reduce( + (text: string, currCell) => text + '\n' + currCell.model.value, + '' + ); + console.log({ text }); + + if (!text) { + return null; + } + + return { + type: 'cell', + widgetId: widget.id, + text + }; +} + +/** + * Gets a Selection object from a document widget. Returns `null` if unable. + */ +export function getSelection(widget: Widget | null): Selection | null { + if (!(widget instanceof DocumentWidget)) { + return null; + } + + const textSelection = getTextSelection(widget); + if (textSelection) { + return textSelection; + } + + const cellSelection = getCellSelection(widget); + if (cellSelection) { + return cellSelection; + } + + return null; +} diff --git a/packages/jupyter-ai/src/selection-watcher.ts b/packages/jupyter-ai/src/selection-watcher.ts index f0e61cf0d..263c7d0be 100644 --- a/packages/jupyter-ai/src/selection-watcher.ts +++ b/packages/jupyter-ai/src/selection-watcher.ts @@ -1,95 +1,13 @@ import { JupyterFrontEnd, LabShell } from '@jupyterlab/application'; import { DocumentWidget } from '@jupyterlab/docregistry'; -import { CodeEditor } from '@jupyterlab/codeeditor'; -import { CodeMirrorEditor } from '@jupyterlab/codemirror'; -import { FileEditor } from '@jupyterlab/fileeditor'; import { Notebook } from '@jupyterlab/notebook'; import { find } from '@lumino/algorithm'; import { Widget } from '@lumino/widgets'; import { Signal } from '@lumino/signaling'; -import { getCellIndex } from './utils'; - -/** - * Gets the editor instance used by a document widget. Returns `null` if unable. - */ -function getEditor(widget: Widget | null) { - if (!(widget instanceof DocumentWidget)) { - return null; - } - - let editor: CodeEditor.IEditor | undefined; - const { content } = widget; - - if (content instanceof FileEditor) { - editor = content.editor; - } else if (content instanceof Notebook) { - editor = content.activeCell?.editor; - } - - if (!(editor instanceof CodeMirrorEditor)) { - return undefined; - } - - return editor; -} - -/** - * Gets a Selection object from a document widget. Returns `null` if unable. - */ -function getTextSelection(widget: Widget | null): Selection | null { - const editor = getEditor(widget); - // widget type check is redundant but hints the type to TypeScript - if (!editor || !(widget instanceof DocumentWidget)) { - return null; - } - - let cellId: string | undefined = undefined; - if (widget.content instanceof Notebook) { - cellId = widget.content.activeCell?.model.id; - } - - let { start, end, ...selectionObj } = editor.getSelection(); - const startOffset = editor.getOffsetAt(start); - const endOffset = editor.getOffsetAt(end); - const text = editor.model.sharedModel - .getSource() - .substring(startOffset, endOffset); - - // ensure start <= end - // required for editor.model.sharedModel.updateSource() - if (startOffset > endOffset) { - [start, end] = [end, start]; - } - - return { - ...selectionObj, - start, - end, - text, - widgetId: widget.id, - ...(cellId && { - cellId - }) - }; -} - -export type Selection = CodeEditor.ITextSelection & { - /** - * The text within the selection as a string. - */ - text: string; - /** - * The ID of the document widget in which the selection was made. - */ - widgetId: string; - /** - * The ID of the cell in which the selection was made, if the original widget - * was a notebook. - */ - cellId?: string; -}; +import type { Selection, TextSelection } from './selection-utils'; +import { getCellIndex, getEditor, getSelection } from './selection-utils'; export class SelectionWatcher { constructor(shell: JupyterFrontEnd.IShell) { @@ -109,7 +27,7 @@ export class SelectionWatcher { return this._selectionChanged; } - replaceSelection(selection: Selection) { + replaceSelection(selection: TextSelection) { // unfortunately shell.currentWidget doesn't update synchronously after // shell.activateById(), which is why we have to get a reference to the // widget manually. @@ -151,7 +69,7 @@ export class SelectionWatcher { protected _poll() { const prevSelection = this._selection; - const currSelection = getTextSelection(this._mainAreaWidget); + const currSelection = getSelection(this._mainAreaWidget); if (prevSelection?.text === currSelection?.text) { return; diff --git a/packages/jupyter-ai/src/utils.ts b/packages/jupyter-ai/src/utils.ts deleted file mode 100644 index 070c5d403..000000000 --- a/packages/jupyter-ai/src/utils.ts +++ /dev/null @@ -1,48 +0,0 @@ -/** - * Contains various utility functions shared throughout the project. - */ -import { Notebook } from '@jupyterlab/notebook'; -import { FileEditor } from '@jupyterlab/fileeditor'; -import { CodeEditor } from '@jupyterlab/codeeditor'; -import { Widget } from '@lumino/widgets'; - -/** - * Get text selection from an editor widget (DocumentWidget#content). - */ -export function getTextSelection(widget: Widget): string { - const editor = getEditor(widget); - if (!editor) { - return ''; - } - - const selectionObj = editor.getSelection(); - const start = editor.getOffsetAt(selectionObj.start); - const end = editor.getOffsetAt(selectionObj.end); - const text = editor.model.sharedModel.getSource().substring(start, end); - - return text; -} - -/** - * Get editor instance from an editor widget (i.e. `DocumentWidget#content`). - */ -export function getEditor(widget: Widget): CodeEditor.IEditor | undefined { - let editor: CodeEditor.IEditor | undefined; - if (widget instanceof FileEditor) { - editor = widget.editor; - } else if (widget instanceof Notebook) { - editor = widget.activeCell?.editor; - } - - return editor; -} - -/** - * Gets the index of the cell associated with `cellId`. - */ -export function getCellIndex(notebook: Notebook, cellId: string): number { - const idx = notebook.model?.sharedModel.cells.findIndex( - cell => cell.getId() === cellId - ); - return idx === undefined ? -1 : idx; -}