Skip to content

Commit

Permalink
mitosheet: automatically import mitosheet package
Browse files Browse the repository at this point in the history
  • Loading branch information
@aarondr77
aarondr77 committed Oct 1, 2024
1 parent 2ea1dfd commit 0e91ff1
Show file tree
Hide file tree
Showing 6 changed files with 106 additions and 55 deletions.
3 changes: 3 additions & 0 deletions mitosheet/dev/macsetup.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,5 +28,8 @@ jlpm run build
# Setup JupyterLab development
jupyter labextension develop . --overwrite

# Enable the server extension
jupyter server extension enable --py mitosheet

# Finally, start watching the javascript
jlpm run watch
7 changes: 7 additions & 0 deletions mitosheet/jupyter-config/jupyter_server_config.d/mitosheet.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
{
"ServerApp": {
"jpserver_extensions": {
"mitosheet": true
}
}
}
30 changes: 20 additions & 10 deletions mitosheet/mitosheet/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@

import os
import pandas as pd
from IPython import get_ipython

# Public interface we want users to rely on
from mitosheet.mito_backend import sheet
Expand Down Expand Up @@ -83,6 +84,24 @@ def _jupyter_nbextension_paths():
}]


def _jupyter_server_extension_points():
"""
Returns a list of dictionaries with metadata describing
where to find the `_load_jupyter_server_extension` function.
"""
return [{"module": "mitosheet"}]

def _load_jupyter_server_extension(server_app):
"""Registers the custom DataFrame formatter when the server extension is loaded."""
import pandas as pd
import mitosheet

print("LOADING!!!!!!!!!!!")

# Log a message to confirm the extension is loaded
server_app.log.info("Mitosheet server extension loaded.")


"""
Dash Integration below.
Expand Down Expand Up @@ -112,16 +131,8 @@ def activate():

# 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 mitosheet.sheet(obj) # Return HTML string
return None # Let other types use the default formatter

def mitosheet_plain_formatter(obj, p, cycle):
Expand All @@ -141,5 +152,4 @@ def mitosheet_plain_formatter(obj, p, cycle):
html_formatter.for_type(pd.DataFrame, mitosheet_display_formatter)
plain_formatter.for_type(pd.DataFrame, mitosheet_plain_formatter)


activate()
12 changes: 12 additions & 0 deletions mitosheet/setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,12 @@ def get_data_files_from_data_files_spec(
return data_files

data_files = get_data_files_from_data_files_spec(data_files_spec)
# Add the jupyter server config file so that the extension is automatically loaded
data_files.append(("etc/jupyter/jupyter_server_config.d", ["jupyter-config/jupyter_server_config.d/mitosheet.json"]))

print("DATA FILES!")
print(data_files)


setup_args = dict(
name = 'mitosheet',
Expand Down Expand Up @@ -84,6 +90,12 @@ def get_data_files_from_data_files_spec(
include_package_data = True,
package_data = {'': ['*.js', '*.css', '*.html']},
data_files = data_files,
entry_points={
"jupyter_serverproxy_servers": [
"mitosheet = mitosheet:_load_jupyter_server_extension",
],
},
jupyter_server_extension="mitosheet", # Automatically enable the server extension
install_requires=[
"jupyterlab~=4.0",
# We allow users to have many versions of pandas installed. All functionality should
Expand Down
6 changes: 0 additions & 6 deletions mitosheet/src/jupyter/extensionUtils.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,24 +39,18 @@ export function getCellIndexByID(cells: CellList | undefined, cellID: string | u
}

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

Check warning on line 51 in mitosheet/src/jupyter/extensionUtils.tsx

View workflow job for this annotation

GitHub Actions / Run linters 

Unexpected any. Specify a different type
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
}
}
Expand Down
103 changes: 64 additions & 39 deletions mitosheet/src/plugin.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@ import {
getCellAtIndex,
getCellCallingMitoshetWithAnalysis,
getCellIndexByExecutionCount,
getCellIndexByID,
getCellText,
getMostLikelyMitosheetCallingCell,
getParentMitoContainer,
Expand All @@ -25,7 +24,7 @@ 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 { KernelConnection } from '@jupyterlab/services';

const registerMitosheetToolbarButtonAdder = (tracker: INotebookTracker) => {

Expand Down Expand Up @@ -53,14 +52,14 @@ const registerMitosheetToolbarButtonAdder = (tracker: INotebookTracker) => {
*/
function activateMitosheetExtension(
app: JupyterFrontEnd,
tracker: INotebookTracker,
notebookTracker: INotebookTracker,
rendermimeRegistry: IRenderMimeRegistry,

Check failure on line 56 in mitosheet/src/plugin.tsx

View workflow job for this annotation

GitHub Actions / Run linters 

'rendermimeRegistry' is defined but never used
): void {

console.log('Mitosheet extension activated');

// Add the Create New Mitosheet button
registerMitosheetToolbarButtonAdder(tracker);
registerMitosheetToolbarButtonAdder(notebookTracker);

/**
* This command creates a new comm for the mitosheet to talk to the mito backend.
Expand All @@ -78,10 +77,10 @@ function activateMitosheetExtension(
// If we get a dummy kernel ID, we are in a jupyterlite instance, so we
// just get the current kernel as this is our best guess as to what the
// kernel is
currentKernel = tracker.currentWidget?.context.sessionContext.session?.kernel;
currentKernel = notebookTracker.currentWidget?.context.sessionContext.session?.kernel;
} else {
// If we have a kernel ID, get the kernel with the correct kernel id
const currentNotebook = tracker.find((nb) => {
const currentNotebook = notebookTracker.find((nb) => {
return nb.sessionContext.session?.kernel?.id === kernelID
});
currentKernel = currentNotebook?.sessionContext?.session?.kernel;
Expand All @@ -106,7 +105,7 @@ function activateMitosheetExtension(
execute: (args: any) => {
const analysisName = args.analysisName as string;
const mitoAPI = args.mitoAPI as MitoAPI;
const cellAndIndex = getMostLikelyMitosheetCallingCell(tracker, analysisName);
const cellAndIndex = getMostLikelyMitosheetCallingCell(notebookTracker, analysisName);

if (cellAndIndex) {
const [cell, ] = cellAndIndex;
Expand All @@ -129,7 +128,7 @@ function activateMitosheetExtension(
const newAnalysisName = args.newAnalysisName as string;
const mitoAPI = args.mitoAPI as MitoAPI;

const mitosheetCallCellAndIndex = getCellCallingMitoshetWithAnalysis(tracker, oldAnalysisName);
const mitosheetCallCellAndIndex = getCellCallingMitoshetWithAnalysis(notebookTracker, oldAnalysisName);
if (mitosheetCallCellAndIndex === undefined) {
return;
}
Expand Down Expand Up @@ -161,14 +160,14 @@ function activateMitosheetExtension(
const code = getCodeString(analysisName, codeLines, telemetryEnabled, publicInterfaceVersion);
// Find the cell that made the mitosheet.sheet call, and if it does not exist, give
// up immediately
const mitosheetCallCellAndIndex = getCellCallingMitoshetWithAnalysis(tracker, analysisName);
const mitosheetCallCellAndIndex = getCellCallingMitoshetWithAnalysis(notebookTracker, analysisName);
if (mitosheetCallCellAndIndex === undefined) {
return;
}

const [, mitosheetCallIndex] = mitosheetCallCellAndIndex;

const notebook = tracker.currentWidget?.content;
const notebook = notebookTracker.currentWidget?.content;
const cells = notebook?.model?.cells;

if (notebook === undefined || cells === undefined) {
Expand Down Expand Up @@ -229,7 +228,7 @@ function activateMitosheetExtension(
// the code in the cell. If they have, we don't want to overwrite their changes automatically.
const oldCode = args.oldCode as string[];
const newCode = getCodeString(analysisName, codeLines, telemetryEnabled, publicInterfaceVersion, true);
const notebook = tracker.currentWidget?.content;
const notebook = notebookTracker.currentWidget?.content;
const cells = notebook?.model?.cells;

if (inputCellExecutionCount === undefined || notebook === undefined || cells === undefined) {
Expand Down Expand Up @@ -294,14 +293,14 @@ function activateMitosheetExtension(
const code = args.code as string;

// Find the cell that made the mitosheet.sheet call, and if it does not exist, give up immediately
const mitosheetCallCellAndIndex = getCellCallingMitoshetWithAnalysis(tracker, analysisName);
const mitosheetCallCellAndIndex = getCellCallingMitoshetWithAnalysis(notebookTracker, analysisName);
if (mitosheetCallCellAndIndex === undefined) {
return;
}

const [, mitosheetCallIndex] = mitosheetCallCellAndIndex;

const notebook = tracker.currentWidget?.content;
const notebook = notebookTracker.currentWidget?.content;
const cells = notebook?.model?.cells;

if (notebook === undefined || cells === undefined) {
Expand Down Expand Up @@ -329,7 +328,7 @@ function activateMitosheetExtension(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
execute: (args: any): string[] => {
const analysisToReplayName = args.analysisToReplayName as string | undefined;
const cellAndIndex = getMostLikelyMitosheetCallingCell(tracker, analysisToReplayName);
const cellAndIndex = getMostLikelyMitosheetCallingCell(notebookTracker, analysisToReplayName);
if (cellAndIndex) {
const [cell, ] = cellAndIndex;
return getArgsFromMitosheetCallCode(getCellText(cell));
Expand All @@ -343,12 +342,12 @@ function activateMitosheetExtension(
label: 'Reads the arguments on the last line of a code cell.',
// eslint-disable-next-line @typescript-eslint/no-explicit-any
execute: (args: any): string[] => {
const notebook = tracker.currentWidget?.content;
const notebook = notebookTracker.currentWidget?.content;
const cells = notebook?.model?.cells;
const cellID = '123' // TODO fix me!
const cellIndex = getCellIndexByID(cells, cellID);
const inputCellExecutionCount = args.inputCellExecutionCount as number | undefined;
const cellIndex = getCellIndexByExecutionCount(cells, inputCellExecutionCount);

if (cellID === undefined || notebook === undefined || cells === undefined || cellIndex === undefined) {
if (notebook === undefined || cells === undefined || cellIndex === undefined) {
return [];
}

Expand Down Expand Up @@ -378,8 +377,8 @@ function activateMitosheetExtension(
console.log("creating mitosheet from dataframe ouput")

// We get the current notebook (currentWidget)
const notebook = tracker.currentWidget?.content;
const context = tracker.currentWidget?.context;
const notebook = notebookTracker.currentWidget?.content;
const context = notebookTracker.currentWidget?.context;
if (!notebook || !context) return;

/*
Expand Down Expand Up @@ -423,8 +422,8 @@ function activateMitosheetExtension(
execute: async (): Promise<void> => {

// We get the current notebook (currentWidget)
const notebook = tracker.currentWidget?.content;
const context = tracker.currentWidget?.context;
const notebook = notebookTracker.currentWidget?.content;
const context = notebookTracker.currentWidget?.context;
if (!notebook || !context) return;

// Create a new code cell that creates a blank mitosheet
Expand Down Expand Up @@ -530,26 +529,52 @@ function activateMitosheetExtension(
selector: '.mito-container'
});


/*
Add a custom mime renderer to handle dataframe outputs.
The section below is responsible for importing the mitosheet package into each
kernel by default. This allows us to render a mitosheet as the default dataframe
renderer even if the user has not yet ran `import mitosheet`
*/

// 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
const importMitosheetPackage = (kernel: KernelConnection | null | undefined) => {
if (kernel) {
// Although I don't think necessary, wrap in a try, except statement for extra safety
kernel.requestExecute({ code: `try:
import mitosheet
except:
pass` });
}
}

// Listen for new notebooks
notebookTracker.widgetAdded.connect((sender, notebookPanel) => {

// When the session changes we need to re-import the package. For example, if the user
// restarts a kernel. Also, becuase the extension is loaded before the kernel is created
// this sessionChange approach is required to import mitosheet when first opening a notebook.
notebookPanel.sessionContext.sessionChanged.connect((sessionContext, sessionContextArgs) => {

Check failure on line 554 in mitosheet/src/plugin.tsx

View workflow job for this annotation

GitHub Actions / Run linters 

'sessionContextArgs' is defined but never used
const kernel = sessionContext.session?.kernel as KernelConnection | null;
importMitosheetPackage(kernel);
});

// Inject code into the current kernel
const kernel = notebookPanel.sessionContext?.session?.kernel as KernelConnection | null;
importMitosheetPackage(kernel);
});

// When activating the Mito extension, import the mitosheet package so we can
// render mitosheets as the default dataframe renderer. TODO: Check if this code
// is ever successful. It might be at this stage in the extension activation process
// the kernel is never defined and so we rely on the above checks instead.
const notebookPanel = notebookTracker.currentWidget
const kernel = notebookPanel?.sessionContext?.session?.kernel as KernelConnection | null;
importMitosheetPackage(kernel)


/*
Finally, add the app commands to the window
so we can write to it elsewhere
*/
window.commands = app.commands;
}

const mitosheetJupyterLabPlugin: JupyterFrontEndPlugin<void> = {
Expand Down

0 comments on commit 0e91ff1

Please sign in to comment.