Fix losses when switching betwee markdown and code (microsoft#8221)

* Fix losses when switching betwee markdown and code

* Fix linting errors

news/2 Fixes/8141.md

@@ -0,0 +1 @@
+Fix for code disappearing after switching between markdown and code in a Notebook Editor.

src/datascience-ui/interactive-common/mainStateController.ts

@@ -565,7 +565,11 @@ export class MainStateController implements IMessageHandler {
         const index = this.pendingState.cellVMs.findIndex(c => c.cell.id === cellId);
         if (index >= 0 && this.pendingState.cellVMs[index].cell.data.cell_type !== newType) {
             const cellVMs = [...this.pendingState.cellVMs];
-            cellVMs[index] = immutable.updateIn(this.pendingState.cellVMs[index], ['cell', 'data', 'cell_type'], () => newType);
+            const current = this.pendingState.cellVMs[index];
+            const newSource = current.focused ? this.getMonacoEditorContents(cellId) : current.cell.data.source;
+            const newCell = { ...current, inputBlockText: newSource, cell: { ...current.cell, state: CellState.executing, data: { ...current.cell.data, cell_type: newType, source: newSource } } };
+            // tslint:disable-next-line: no-any
+            cellVMs[index] = (newCell as any); // This is because IMessageCell doesn't fit in here. But message cells can't change type
             this.setState({ cellVMs });
             if (newType === 'code') {
                 this.sendMessage(InteractiveWindowMessages.InsertCell, { cell: cellVMs[index].cell, index, code: concatMultilineStringInput(cellVMs[index].cell.data.source), codeCellAboveId: this.firstCodeCellAbove(cellId) });
@@ -806,6 +810,20 @@ export class MainStateController implements IMessageHandler {
         return cellVM;
     }
 
+    protected getMonacoEditorContents(cellId: string): string | undefined {
+        const index = this.findCellIndex(cellId);
+        if (index >= 0) {
+            // Get the model for the monaco editor
+            const monacoId = this.getMonacoId(cellId);
+            if (monacoId) {
+                const model = monacoEditor.editor.getModels().find(m => m.id === monacoId);
+                if (model) {
+                    return model.getValue().replace(/\r/g, '');
+                }
+            }
+        }
+    }
+
     protected onCodeLostFocus(_cellId: string) {
         // Default is do nothing.
     }

src/datascience-ui/native-editor/nativeCell.tsx

@@ -574,35 +574,13 @@ export class NativeCell extends React.Component<INativeCellProps> {
         const canRunBelow = this.props.cellVM.cell.state === CellState.finished || this.props.cellVM.cell.state === CellState.error;
         const switchTooltip = this.props.cellVM.cell.data.cell_type === 'code' ? getLocString('DataScience.switchToMarkdown', 'Change to markdown') :
             getLocString('DataScience.switchToCode', 'Change to code');
-        const switchToMarkdown = () => {
-            this.props.stateController.changeCellType(cellId, 'markdown');
-            this.props.stateController.sendCommand(NativeCommandType.ChangeToMarkdown, 'mouse');
-            setTimeout(() => this.props.focusCell(cellId, true), 0);
+        const otherCellType = this.props.cellVM.cell.data.cell_type === 'code' ? 'markdown' : 'code';
+        const otherCellTypeCommand = otherCellType === 'markdown' ? NativeCommandType.ChangeToMarkdown : NativeCommandType.ChangeToCode;
+        const otherCellImage = otherCellType === 'markdown' ? ImageName.SwitchToMarkdown : ImageName.SwitchToCode;
+        const switchCellType = () => {
+            this.props.stateController.changeCellType(cellId, otherCellType);
+            this.props.stateController.sendCommand(otherCellTypeCommand, 'mouse');
         };
-        const switchToCode = () => {
-            const handler = () => {
-                setTimeout(() => {
-                    this.props.stateController.changeCellType(cellId, 'code');
-                    this.props.stateController.sendCommand(NativeCommandType.ChangeToCode, 'mouse');
-                    this.props.focusCell(cellId, true);
-                }, 0);
-            };
-
-            // This is special. Coming in on a mouse down event so we get
-            // called before focus changes. After focus changes, then switch to code
-            if (this.isShowingMarkdownEditor()) {
-                this.pendingFocusLoss = handler;
-            } else {
-                handler();
-            }
-        };
-        const switchButton = this.props.cellVM.cell.data.cell_type === 'code' ?
-            <ImageButton baseTheme={this.props.baseTheme} onClick={switchToMarkdown} tooltip={switchTooltip}>
-                <Image baseTheme={this.props.baseTheme} class='image-button-image' image={ImageName.SwitchToMarkdown} />
-            </ImageButton> :
-            <ImageButton baseTheme={this.props.baseTheme} onMouseDown={switchToCode} tooltip={switchTooltip}>
-                <Image baseTheme={this.props.baseTheme} class='image-button-image' image={ImageName.SwitchToCode} />
-            </ImageButton>;
 
         return (
             <div className='native-editor-celltoolbar-middle'>
@@ -612,7 +590,9 @@ export class NativeCell extends React.Component<INativeCellProps> {
                 <ImageButton baseTheme={this.props.baseTheme} onClick={runBelow} disabled={!canRunBelow} tooltip={getLocString('DataScience.runBelow', 'Run cell and below')}>
                     <Image baseTheme={this.props.baseTheme} class='image-button-image' image={ImageName.RunBelow} />
                 </ImageButton>
-                {switchButton}
+                <ImageButton baseTheme={this.props.baseTheme} onMouseDown={switchCellType} tooltip={switchTooltip}>
+                    <Image baseTheme={this.props.baseTheme} class='image-button-image' image={otherCellImage} />
+                </ImageButton>
                 <ImageButton baseTheme={this.props.baseTheme} onClick={deleteCell} tooltip={getLocString('DataScience.deleteCell', 'Delete cell')}>
                     <Image baseTheme={this.props.baseTheme} class='image-button-image' image={ImageName.Delete} />
                 </ImageButton>

src/datascience-ui/native-editor/nativeEditorStateController.ts

@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 'use strict';
-import * as monacoEditor from 'monaco-editor/esm/vs/editor/editor.api';
 import * as uuid from 'uuid/v4';
 
 import { noop } from '../../client/common/utils/misc';
@@ -285,28 +284,25 @@ export class NativeEditorStateController extends MainStateController {
         // Update the cell's source
         const index = this.findCellIndex(cellId);
         if (index >= 0) {
-            // Get the model for the monaco editor
-            const monacoId = this.getMonacoId(cellId);
-            if (monacoId) {
-                const model = monacoEditor.editor.getModels().find(m => m.id === monacoId);
-                if (model) {
-                    const newValue = model.getValue().replace(/\r/g, '');
-                    const newVMs = [...this.getState().cellVMs];
-
-                    // Update our state
-                    newVMs[index] = {
-                        ...newVMs[index],
-                        cell: {
-                            ...newVMs[index].cell,
-                            data: {
-                                ...newVMs[index].cell.data,
-                                source: newValue
-                            }
+            // Get the model source from the monaco editor
+            const source = this.getMonacoEditorContents(cellId);
+            if (source) {
+                const newVMs = [...this.getState().cellVMs];
+
+                // Update our state
+                newVMs[index] = {
+                    ...newVMs[index],
+                    inputBlockText: source,
+                    cell: {
+                        ...newVMs[index].cell,
+                        data: {
+                            ...newVMs[index].cell.data,
+                            source
                         }
-                    };
+                    }
+                };
 
-                    this.setState({ cellVMs: newVMs });
-                }
+                this.setState({ cellVMs: newVMs });
             }
         }
     }

src/test/datascience/nativeEditor.functional.test.tsx

@@ -38,6 +38,7 @@ import {
     addCell,
     closeNotebook,
     createNewEditor,
+    focusCell,
     getNativeCellResults,
     mountNativeWebView,
     openEditor,
@@ -55,6 +56,7 @@ import {
     getLastOutputCell,
     getNativeFocusedEditor,
     getOutputCell,
+    injectCode,
     isCellFocused,
     isCellSelected,
     srcDirectory,
@@ -891,6 +893,14 @@ for _ in range(50):
                     0
                 );
 
+                // Force focus so we can change the text. Use special method
+                // because we can't key down on the editor
+                await focusCell(ioc, wrapper, 1);
+
+                // Change the markdown
+                let editor = getNativeFocusedEditor(wrapper);
+                injectCode(editor, 'foo');
+
                 // Switch back to code mode.
                 // At this moment, there's no cell input element, hence send key strokes to the wrapper.
                 const wrapperElement = wrapper
@@ -908,6 +918,11 @@ for _ in range(50):
                         .find(MonacoEditor).length,
                     1
                 );
+
+                // Confirm editor still has the same text
+                editor = getNativeFocusedEditor(wrapper);
+                const monacoEditor = editor!.instance() as MonacoEditor;
+                assert.equal('foo', monacoEditor.state.editor!.getValue(), 'Changing cell type lost input');
             });
 
             test('Test undo using the key \'z\'', async () => {

src/test/datascience/nativeEditorTestHelpers.tsx

@@ -8,6 +8,7 @@ import { Uri } from 'vscode';
 
 import { InteractiveWindowMessages } from '../../client/datascience/interactive-common/interactiveWindowTypes';
 import { IJupyterExecution, INotebookEditor, INotebookEditorProvider } from '../../client/datascience/types';
+import { NativeCell } from '../../datascience-ui/native-editor/nativeCell';
 import { NativeEditor } from '../../datascience-ui/native-editor/nativeEditor';
 import { ImageButton } from '../../datascience-ui/react-common/imageButton';
 import { DataScienceIocContainer } from './dataScienceIocContainer';
@@ -79,6 +80,16 @@ export async function setupWebview(ioc: DataScienceIocContainer) {
     }
 }
 
+export function focusCell(ioc: DataScienceIocContainer, wrapper: ReactWrapper<any, Readonly<{}>, React.Component>, index: number): Promise<void> {
+    const focusChange = waitForMessage(ioc, InteractiveWindowMessages.FocusedCellEditor);
+    const cell = wrapper.find(NativeCell).at(index);
+    if (cell) {
+        const vm = cell.props().cellVM;
+        cell.props().focusCell(vm.cell.id, true);
+    }
+    return focusChange;
+}
+
 // tslint:disable-next-line: no-any
 export async function addCell(wrapper: ReactWrapper<any, Readonly<{}>, React.Component>, ioc: DataScienceIocContainer, code: string, submit: boolean = true): Promise<void> {
     // First get the main toolbar. We'll use this to add a cell.

src/test/datascience/testHelpers.tsx

@@ -420,6 +420,13 @@ export function getInteractiveEditor(wrapper: ReactWrapper<any, Readonly<{}>, Re
     return lastCell.find('MonacoEditor');
 }
 
+export function getNativeEditor(wrapper: ReactWrapper<any, Readonly<{}>, React.Component>, index: number): ReactWrapper<any, Readonly<{}>, React.Component> | undefined {
+    // Find the last cell. It should have a monacoEditor object
+    const cells = wrapper.find('NativeCell');
+    const lastCell = index < cells.length ? cells.at(index) : undefined;
+    return lastCell ? lastCell.find('MonacoEditor') : undefined;
+}
+
 export function getNativeFocusedEditor(wrapper: ReactWrapper<any, Readonly<{}>, React.Component>): ReactWrapper<any, Readonly<{}>, React.Component> | undefined {
     // Find the last cell. It should have a monacoEditor object
     wrapper.update();