mitosheet: new approach -- overwriting dataframe ipython render

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'}, 
-]
+]
+
+
+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()

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(
         <script>
             {js_code}
         </script>
-    </div>""")) # type: ignore
+    </div>""")) # type: ignore
+
+

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,

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();
     }
 }

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 ''; 
 

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<string[]> => {
+export const getArgs = (analysisToReplayName: string | undefined, inputCellExecutionCount: number | undefined): Promise<string[]> => {
     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 {

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, 

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, 

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<string[]>,
+    getArgs: (analysisToReplayName: string | undefined, inputCellExecutionCount: number | undefined) => Promise<string[]>,
     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
 }