diff --git a/mitosheet/mitosheet/__init__.py b/mitosheet/mitosheet/__init__.py index a642cdd4d..65ee69f18 100644 --- a/mitosheet/mitosheet/__init__.py +++ b/mitosheet/mitosheet/__init__.py @@ -102,4 +102,44 @@ def _jupyter_nbextension_paths(): _css_dist = [ {'relative_package_path': 'mito_dash/v1/mitoBuild/component.css', 'namespace': 'mitosheet'}, -] \ No newline at end of file +] + + +def activate(): + from IPython import get_ipython + import pandas as pd + import mitosheet + + # Updated formatter functions with correct signatures + def mitosheet_display_formatter(obj, include=None, exclude=None): + + # We do not have access to the cell ID here because the cell ID exists only in the frontend + # and does not get shared with the kernel. However, we do get access to the execution count, + # which is also a unique identifier for each cell within the lifecycle of each kernel. + ip = get_ipython() + print('ip', ip.execution_count) + + + if isinstance(obj, pd.DataFrame): + return mitosheet.sheet(obj, input_cell_execution_count = ip.execution_count) # Return HTML string + return None # Let other types use the default formatter + + def mitosheet_plain_formatter(obj, p, cycle): + if isinstance(obj, pd.DataFrame): + return '' # Prevent default text representation + return None # Let other types use the default formatter + + ip = get_ipython() + html_formatter = ip.display_formatter.formatters['text/html'] + plain_formatter = ip.display_formatter.formatters['text/plain'] + + # Save the original formatters + activate.original_html_formatter = html_formatter.for_type(pd.DataFrame) + activate.original_plain_formatter = plain_formatter.for_type(pd.DataFrame) + + # Register the custom formatters + html_formatter.for_type(pd.DataFrame, mitosheet_display_formatter) + plain_formatter.for_type(pd.DataFrame, mitosheet_plain_formatter) + + +activate() \ No newline at end of file diff --git a/mitosheet/mitosheet/mito_backend.py b/mitosheet/mitosheet/mito_backend.py index 74ac34fdd..a32c88792 100644 --- a/mitosheet/mitosheet/mito_backend.py +++ b/mitosheet/mitosheet/mito_backend.py @@ -60,7 +60,7 @@ def __init__( column_definitions: Optional[List[ColumnDefinitions]]=None, default_editing_mode: Optional[DefaultEditingMode]=None, theme: Optional[MitoTheme]=None, - cell_id: Optional[str]=None, + input_cell_execution_count: Optional[int]=None, ): """ Takes a list of dataframes and strings that are paths to CSV files @@ -110,7 +110,7 @@ def __init__( column_definitions=column_definitions, theme=theme, default_editing_mode=default_editing_mode, - cell_id=cell_id + input_cell_execution_count=input_cell_execution_count ) # And the api @@ -321,7 +321,7 @@ def get_mito_backend( user_defined_importers: Optional[List[Callable]]=None, user_defined_editors: Optional[List[Callable]]=None, column_definitions: Optional[List[ColumnDefinitions]]=None, - cell_id: Optional[str]=None, + input_cell_execution_count: Optional[int]=None, ) -> MitoBackend: # We pass in the dataframes directly to the widget @@ -332,7 +332,7 @@ def get_mito_backend( user_defined_importers=user_defined_importers, user_defined_editors=user_defined_editors, column_definitions=column_definitions, - cell_id=cell_id + input_cell_execution_count=input_cell_execution_count ) return mito_backend @@ -393,7 +393,7 @@ def sheet( sheet_functions: Optional[List[Callable]]=None, importers: Optional[List[Callable]]=None, editors: Optional[List[Callable]]=None, - cell_id: Optional[str]=None # If the sheet is a dataframe mime renderer, we pass the cell_id so we know where to generate the code. + input_cell_execution_count: Optional[int]=None # If the sheet is a dataframe mime renderer, we pass the cell_id so we know where to generate the code. ) -> None: """ Renders a Mito sheet. If no arguments are passed, renders an empty sheet. Otherwise, renders @@ -446,7 +446,7 @@ def sheet( user_defined_functions=sheet_functions, user_defined_importers=importers, user_defined_editors=editors, - cell_id=cell_id + input_cell_execution_count=input_cell_execution_count ) # Setup the comm target on this @@ -482,4 +482,6 @@ def sheet( - """)) # type: ignore \ No newline at end of file + """)) # type: ignore + + diff --git a/mitosheet/mitosheet/steps_manager.py b/mitosheet/mitosheet/steps_manager.py index 1e21041f2..4e658ca26 100644 --- a/mitosheet/mitosheet/steps_manager.py +++ b/mitosheet/mitosheet/steps_manager.py @@ -190,7 +190,7 @@ def __init__( column_definitions: Optional[List[ColumnDefinitions]]=None, default_editing_mode: Optional[DefaultEditingMode]=None, theme: Optional[MitoTheme]=None, - cell_id: Optional[str]=None, + input_cell_execution_count: Optional[int]=None, ): """ When initalizing the StepsManager, we also do preprocessing @@ -207,7 +207,7 @@ def __init__( # is such an analysis self.analysis_to_replay = analysis_to_replay self.analysis_to_replay_exists = get_analysis_exists(analysis_to_replay) - self.cell_id = cell_id + self.input_cell_execution_count = input_cell_execution_count # The import folder is the folder that users have the right to import files from. # If this is set, then we should never let users view or access files that are not @@ -405,7 +405,7 @@ def analysis_data_json(self): return json.dumps( { "analysisName": self.analysis_name, - "cellID": self.cell_id, + "inputCellExecutionCount": self.input_cell_execution_count, "publicInterfaceVersion": self.public_interface_version, "analysisToReplay": { 'analysisName': self.analysis_to_replay, diff --git a/mitosheet/src/DataFrameMimeRenderer.tsx b/mitosheet/src/DataFrameMimeRenderer.tsx index 856de3aa1..e048b91c6 100644 --- a/mitosheet/src/DataFrameMimeRenderer.tsx +++ b/mitosheet/src/DataFrameMimeRenderer.tsx @@ -8,8 +8,6 @@ import { getLastNonEmptyLine } from './jupyter/code'; import { CodeCell } from '@jupyterlab/cells'; -const CLASS_NAME = 'jp-DataFrameViewer'; - const SpreadsheetDataframeComponent = (props: { htmlContent: string, jsCode?: string }) => { /** * The `useEffect` hook is used here to ensure that the JavaScript code is executed @@ -51,7 +49,6 @@ export class DataFrameMimeRenderer extends Widget implements IRenderMime.IRender defaultRenderer: IRenderMime.IRenderer ) { super(); - this.addClass(CLASS_NAME); this._notebookTracker = notebookTracker; this._defaultRenderer = defaultRenderer;; } @@ -84,10 +81,18 @@ export class DataFrameMimeRenderer extends Widget implements IRenderMime.IRender 3. Get the code cell ID from the code cell's model. 4. Use the cell ID to find the input cell and read the dataframe name from it. + + TODO: There is still a race condition bug where if code cell 1 creates a dataframe renderer, and code cell 2 + edits the dataframe, the mitosheet output will show the dataframe state after code cell 2 has run, instead of the + state of the dataframe at code cell 1. :/ + + One idea, compare the dataframe we render and the dataframe in the default mime renderer. If they are different, + default to showing the default dataframe output and add a message that says "rerun the code cell above to see + the current dataframe in mito" + */ try { - const isDataframeOutput = model.data['text/html']?.toString()?.includes('class="dataframe"'); if (!isDataframeOutput) { @@ -146,8 +151,14 @@ export class DataFrameMimeRenderer extends Widget implements IRenderMime.IRender throw new Error('No kernel found'); } - const pythonCode = `import mitosheet; mitosheet.sheet(${dataframeVariableName || ''}, cell_id='${inputCellID}')`; - const future = kernel.requestExecute({ code: pythonCode }); + + const newDataframeName = `${dataframeVariableName}_mito` + + const pythonCode = `import mitosheet; ${newDataframeName} = ${dataframeVariableName}.copy(deep=True);mitosheet.sheet(${newDataframeName}, cell_id='${inputCellID}')`; + const future = kernel.requestExecute({ code: pythonCode, silent: true }); + + + console.log('has pending input', kernel.hasPendingInput) /* Listen to the juptyer messages to find the response to the mitosheet.sheet() call. Once we get back the @@ -171,6 +182,9 @@ export class DataFrameMimeRenderer extends Widget implements IRenderMime.IRender // Attatch the Mito widget to the node this.node.innerHTML = ''; Widget.attach(reactWidget, this.node); + console.log('attatched!!!!!!!!!!!!!!!!!!!!!!!!!!!!!') + } else { + console.log('msg', msg) } }; } catch (error) { @@ -180,6 +194,7 @@ export class DataFrameMimeRenderer extends Widget implements IRenderMime.IRender this.node.replaceWith(this._defaultRenderer.node); } + console.log('after!!!!!!!!!!!!!!!!!!!!!!!!!') return Promise.resolve(); } } diff --git a/mitosheet/src/jupyter/extensionUtils.tsx b/mitosheet/src/jupyter/extensionUtils.tsx index 057c13fa9..9c48ef44a 100644 --- a/mitosheet/src/jupyter/extensionUtils.tsx +++ b/mitosheet/src/jupyter/extensionUtils.tsx @@ -38,6 +38,49 @@ export function getCellIndexByID(cells: CellList | undefined, cellID: string | u return cellIndex === -1 ? undefined : cellIndex; } +export function getCellIndexByExecutionCount(cells: CellList | undefined, executionCount: number | undefined): number | undefined { + console.log("starting function") + if (cells == undefined || executionCount == undefined) { + return undefined; + } + + console.log("finding cell index by execution count", executionCount) + + // In order to get the cell index, we need to iterate over the cells and call the `get` method + // to see the cells in order. Otherwise, the cells are returned in a random order. + for (let i = 0; i < cells.length; i++) { + const cell = cells.get(i) + console.log('cell', cell) + if (cell.type === 'code') { + const nonTypeSafeCell = cell as any + const executionCountEntry = nonTypeSafeCell.sharedModel.ymodel._map.get('execution_count')['content']['arr'][0]; + console.log('executionCountEntry', executionCountEntry) + if (executionCountEntry === executionCount) { + console.log("returning", i) + return i + } + } + } + + return undefined +} + +export function getCellByExecutionCount(cells: CellList | undefined, executionCount: number | undefined): ICellModel | undefined { + if (cells == undefined || executionCount == undefined) { + return undefined; + } + + return Array.from(cells).find((cell: ICellModel) => { + if (cell.type === 'code') { + const nonTypeSafeCell = cell as any + const executionCountEntry = nonTypeSafeCell.sharedModel.ymodel._map.get('execution_count'); + if (executionCountEntry == executionCount) { + return cell + } + } + }) +} + export function getCellText(cell: ICellModel| undefined): string { if (cell == undefined) return ''; diff --git a/mitosheet/src/jupyter/jupyterUtils.tsx b/mitosheet/src/jupyter/jupyterUtils.tsx index 27368e909..56319d326 100644 --- a/mitosheet/src/jupyter/jupyterUtils.tsx +++ b/mitosheet/src/jupyter/jupyterUtils.tsx @@ -47,12 +47,12 @@ export const overwriteAnalysisToReplayToMitosheetCall = (oldAnalysisName: string } -export const writeGeneratedCodeToCell = (analysisName: string, cellID: string | undefined, code: string[], telemetryEnabled: boolean, publicInterfaceVersion: PublicInterfaceVersion, triggerUserEditedCodeDialog: (codeWithoutUserEdits: string[], codeWithUserEdits: string[]) => void, oldCode: string[], overwriteIfUserEditedCode?: boolean): void => { +export const writeGeneratedCodeToCell = (analysisName: string, inputCellExecutionCount: number | undefined, code: string[], telemetryEnabled: boolean, publicInterfaceVersion: PublicInterfaceVersion, triggerUserEditedCodeDialog: (codeWithoutUserEdits: string[], codeWithUserEdits: string[]) => void, oldCode: string[], overwriteIfUserEditedCode?: boolean): void => { if (isInJupyterLabOrNotebook()) { - if (cellID) { + if (inputCellExecutionCount) { window.commands?.execute('mitosheet:write-generated-code-cell-by-id', { analysisName: analysisName, - cellID: cellID, + inputCellExecutionCount: inputCellExecutionCount, code: code, telemetryEnabled: telemetryEnabled, publicInterfaceVersion: publicInterfaceVersion, @@ -86,11 +86,11 @@ export const writeCodeSnippetCell = (analysisName: string, code: string): void = } -export const getArgs = (analysisToReplayName: string | undefined, cellID: string | undefined): Promise => { +export const getArgs = (analysisToReplayName: string | undefined, inputCellExecutionCount: number | undefined): Promise => { return new Promise((resolve) => { if (isInJupyterLabOrNotebook()) { - if (cellID) { - window.commands?.execute('mitosheet:get-args-by-id', {cellID: cellID}).then(async (args: string[]) => { + if (inputCellExecutionCount) { + window.commands?.execute('mitosheet:get-args-by-id', {inputCellExecutionCount: inputCellExecutionCount}).then(async (args: string[]) => { return resolve(args); }) } else { diff --git a/mitosheet/src/mito/Mito.tsx b/mitosheet/src/mito/Mito.tsx index 019219e29..e66bfc82c 100644 --- a/mitosheet/src/mito/Mito.tsx +++ b/mitosheet/src/mito/Mito.tsx @@ -182,7 +182,7 @@ export const Mito = (props: MitoProps): JSX.Element => { const updateMitosheetCallCellOnFirstRender = async () => { // Then, we go and read the arguments to the mitosheet.sheet() call. If there // is an analysis to replay, we use this to help lookup the call - const args = await props.jupyterUtils?.getArgs(analysisData.analysisToReplay?.analysisName, analysisData.cellID) ?? []; + const args = await props.jupyterUtils?.getArgs(analysisData.analysisToReplay?.analysisName, analysisData.inputCellExecutionCount) ?? []; // Then, after we have the args, we replay an analysis if there is an analysis to replay // Note that this has to happen after so that we have the the argument names loaded in at @@ -300,7 +300,7 @@ export const Mito = (props: MitoProps): JSX.Element => { props.jupyterUtils?.writeGeneratedCodeToCell( analysisData.analysisName, - analysisData.cellID, + analysisData.inputCellExecutionCount, analysisData.code, userProfile.telemetryEnabled, analysisData.publicInterfaceVersion, diff --git a/mitosheet/src/mito/components/modals/UserEditedCodeModal.tsx b/mitosheet/src/mito/components/modals/UserEditedCodeModal.tsx index 5f35bf092..bbffdbe18 100644 --- a/mitosheet/src/mito/components/modals/UserEditedCodeModal.tsx +++ b/mitosheet/src/mito/components/modals/UserEditedCodeModal.tsx @@ -36,7 +36,7 @@ const UserEditedCodeModal = ( const handleUserEditedCode = (overwriteCode: boolean) => { props.jupyterUtils?.writeGeneratedCodeToCell( props.analysisData.analysisName, - props.analysisData.cellID, + props.analysisData.inputCellExecutionCount, props.analysisData.code, props.userProfile.telemetryEnabled, props.analysisData.publicInterfaceVersion, diff --git a/mitosheet/src/mito/types.tsx b/mitosheet/src/mito/types.tsx index a1f645dd4..3211f352b 100644 --- a/mitosheet/src/mito/types.tsx +++ b/mitosheet/src/mito/types.tsx @@ -778,7 +778,7 @@ export type UserDefinedFunction = { */ export interface AnalysisData { analysisName: string, - cellID: string | undefined, + inputCellExecutionCount: number | undefined, publicInterfaceVersion: PublicInterfaceVersion, analysisToReplay: { analysisName: string, @@ -997,9 +997,9 @@ export const enum FeedbackID { } export interface JupyterUtils { - getArgs: (analysisToReplayName: string | undefined, cellID: string | undefined) => Promise, + getArgs: (analysisToReplayName: string | undefined, inputCellExecutionCount: number | undefined) => Promise, writeAnalysisToReplayToMitosheetCall: (analysisName: string, mitoAPI: MitoAPI) => void - writeGeneratedCodeToCell: (analysisName: string, cellID: string | undefined, code: string[], telemetryEnabled: boolean, publicInterfaceVersion: PublicInterfaceVersion, triggerUserEditedCodeDialog: (codeWithoutUserEdits: string[], codeWithUserEdits: string[]) => void, oldCode: string[], overwriteIfUserEditedCode?: boolean) => void + writeGeneratedCodeToCell: (analysisName: string, inputCellExecutionCount: number | undefined, code: string[], telemetryEnabled: boolean, publicInterfaceVersion: PublicInterfaceVersion, triggerUserEditedCodeDialog: (codeWithoutUserEdits: string[], codeWithUserEdits: string[]) => void, oldCode: string[], overwriteIfUserEditedCode?: boolean) => void writeCodeSnippetCell: (analysisName: string, code: string) => void overwriteAnalysisToReplayToMitosheetCall: (oldAnalysisName: string, newAnalysisName: string, mitoAPI: MitoAPI) => void } diff --git a/mitosheet/src/plugin.tsx b/mitosheet/src/plugin.tsx index 9ff1f530b..6cfed06bc 100644 --- a/mitosheet/src/plugin.tsx +++ b/mitosheet/src/plugin.tsx @@ -8,16 +8,23 @@ import { getArgsFromMitosheetCallCode, getCodeString, getLastNonEmptyLine, hasCo import { JupyterComm } from './jupyter/comm'; import { createCodeCellAtIndex, - getCellAtIndex, getCellCallingMitoshetWithAnalysis, getCellIndexByID, getCellText, getMostLikelyMitosheetCallingCell, getParentMitoContainer, isEmptyCell, tryOverwriteAnalysisToReplayParameter, tryWriteAnalysisToReplayParameter, writeToCell, + getCellAtIndex, + getCellCallingMitoshetWithAnalysis, + getCellIndexByExecutionCount, + getCellIndexByID, + getCellText, + getMostLikelyMitosheetCallingCell, + getParentMitoContainer, + isEmptyCell, + tryOverwriteAnalysisToReplayParameter, + tryWriteAnalysisToReplayParameter, + writeToCell, writeToCodeCellAtIndex } from './jupyter/extensionUtils'; import { MitoAPI, PublicInterfaceVersion } from './mito'; import { MITO_TOOLBAR_OPEN_SEARCH_ID, MITO_TOOLBAR_REDO_ID, MITO_TOOLBAR_UNDO_ID } from './mito/components/toolbar/Toolbar'; import { getOperatingSystem, keyboardShortcuts } from './mito/utils/keyboardShortcuts'; import { IRenderMimeRegistry} from '@jupyterlab/rendermime'; -import { IRenderMime } from '@jupyterlab/rendermime-interfaces'; -import DataFrameMimeRenderer from './DataFrameMimeRenderer'; -import { CodeCell } from '@jupyterlab/cells'; const registerMitosheetToolbarButtonAdder = (tracker: INotebookTracker) => { @@ -216,7 +223,7 @@ function activateMitosheetExtension( const codeLines = args.code as string[]; const telemetryEnabled = args.telemetryEnabled as boolean; const publicInterfaceVersion = args.publicInterfaceVersion as PublicInterfaceVersion; - const cellID = args.cellID as string | undefined; + const inputCellExecutionCount = args.inputCellExecutionCount as number | undefined; // This is the last saved analysis' code, which we use to check if the user has changed // the code in the cell. If they have, we don't want to overwrite their changes automatically. @@ -225,11 +232,12 @@ function activateMitosheetExtension( const notebook = tracker.currentWidget?.content; const cells = notebook?.model?.cells; - if (cellID === undefined || notebook === undefined || cells === undefined) { + if (inputCellExecutionCount === undefined || notebook === undefined || cells === undefined) { return; } - const mimeRenderInputCellIndex = getCellIndexByID(cells, cellID); + const mimeRenderInputCellIndex = getCellIndexByExecutionCount(cells, inputCellExecutionCount); + console.log('mimeRenderInputCellIndex', mimeRenderInputCellIndex) if (mimeRenderInputCellIndex === undefined) { // If the code cell that created the mitosheet mime render does not exist, // just return. I don't think this should ever happen because you can't @@ -337,7 +345,7 @@ function activateMitosheetExtension( execute: (args: any): string[] => { const notebook = tracker.currentWidget?.content; const cells = notebook?.model?.cells; - const cellID = args.cellID as string | undefined; + const cellID = '123' // TODO fix me! const cellIndex = getCellIndexByID(cells, cellID); if (cellID === undefined || notebook === undefined || cells === undefined || cellIndex === undefined) { @@ -522,44 +530,24 @@ function activateMitosheetExtension( selector: '.mito-container' }); - console.log('originalExecute', CodeCell.prototype) - - const originalExecute = CodeCell.execute; - CodeCell.execute = async function ( - cell: CodeCell, - sessionContext: any, - metadata: Record = {} - ) { - // Ensure metadata exists - metadata = { ...metadata }; - - // Add cell ID to transient metadata - metadata.transient = { - ...(metadata.transient || {}), - cellId: cell.model.id - }; - - // Call the original execute method with updated metadata - const result = await originalExecute.call(this, cell, sessionContext, metadata); - - return result; - }; - - // Add a custom renderer for the stderr output - - const dataframeMimeType = 'text/html' - const factory = rendermimeRegistry.getFactory(dataframeMimeType); - - if (factory) { - rendermimeRegistry.addFactory({ - safe: true, - mimeTypes: [dataframeMimeType], // Include both MIME types as needed - createRenderer: (options: IRenderMime.IRendererOptions) => { - const defaultRenderer = factory.createRenderer(options); - return new DataFrameMimeRenderer(options, tracker, defaultRenderer); // Pass dataframe to your renderer - } - }, -1); // Giving this renderer a lower rank than the default renderer gives this default priority -} + + /* + Add a custom mime renderer to handle dataframe outputs. + */ + + // const dataframeMimeType = 'text/html' + // const factory = rendermimeRegistry.getFactory(dataframeMimeType); + + // if (factory) { + // rendermimeRegistry.addFactory({ + // safe: true, + // mimeTypes: [dataframeMimeType], + // createRenderer: (options: IRenderMime.IRendererOptions) => { + // const defaultRenderer = factory.createRenderer(options); + // return new DataFrameMimeRenderer(options, tracker, defaultRenderer); + // } + // }, -1); // Giving this renderer a lower rank than the default renderer gives this default priority + // } window.commands = app.commands; // So we can write to it elsewhere }