claim ownership of python file if we create code cells for it (#16103)

amunger · web-flow · commit c8287cd5e68a · 2024-10-04T15:43:26.000-07:00
* claim ownership of python file if we create code cells for it

* remove parameter

* update unit tests
diff --git a/src/interactive-window/editor-integration/codelensprovider.ts b/src/interactive-window/editor-integration/codelensprovider.ts
@@ -37,6 +37,7 @@ export class DataScienceCodeLensProvider implements IDataScienceCodeLensProvider
     private totalGetCodeLensCalls: number = 0;
     private activeCodeWatchers: ICodeWatcher[] = [];
     private didChangeCodeLenses: vscode.EventEmitter<void> = new vscode.EventEmitter<void>();
+    private cachedOwnsSetting: boolean;
 
     constructor(
         @inject(IServiceContainer) private serviceContainer: IServiceContainer,
@@ -59,6 +60,10 @@ export class DataScienceCodeLensProvider implements IDataScienceCodeLensProvider
         }
 
         disposableRegistry.push(vscode.window.onDidChangeActiveTextEditor(() => this.onChangedActiveTextEditor()));
+        const settings = this.configuration.getSettings(undefined);
+        this.cachedOwnsSetting = settings.sendSelectionToInteractiveWindow;
+        this.updateOwnerContextKey();
+        disposableRegistry.push(vscode.workspace.onDidChangeConfiguration((e) => this.onSettingChanged(e)));
         this.onChangedActiveTextEditor();
     }
 
@@ -73,9 +78,33 @@ export class DataScienceCodeLensProvider implements IDataScienceCodeLensProvider
             // set the context to false so our command doesn't run for other files
             const hasCellsContext = new ContextKey(EditorContexts.HasCodeCells);
             hasCellsContext.set(false).catch((ex) => logger.warn('Failed to set jupyter.HasCodeCells context', ex));
+            this.updateOwnerContextKey(false);
         }
     }
 
+    private onSettingChanged(e: vscode.ConfigurationChangeEvent) {
+        if (e.affectsConfiguration('jupyter.interactiveWindow.textEditor.executeSelection')) {
+            const settings = this.configuration.getSettings(undefined);
+            this.cachedOwnsSetting = settings.sendSelectionToInteractiveWindow;
+            this.updateOwnerContextKey();
+        }
+    }
+
+    private updateOwnerContextKey(hasCodeCells?: boolean) {
+        const editorContext = new ContextKey(EditorContexts.OwnsSelection);
+        if (this.cachedOwnsSetting) {
+            editorContext.set(true).catch(noop);
+            return;
+        }
+
+        if (hasCodeCells === undefined) {
+            const hasCellsContext = new ContextKey(EditorContexts.HasCodeCells);
+            hasCodeCells = hasCellsContext.value ?? false;
+        }
+
+        editorContext.set(hasCodeCells).catch(noop);
+    }
+
     public dispose() {
         // On shutdown send how long on average we spent parsing code lens
         if (this.totalGetCodeLensCalls > 0) {
@@ -135,9 +164,9 @@ export class DataScienceCodeLensProvider implements IDataScienceCodeLensProvider
         // ask whenever a change occurs. Do this regardless of if we have code lens turned on or not as
         // shift+enter relies on this code context.
         const hasCellsContext = new ContextKey(EditorContexts.HasCodeCells);
-        hasCellsContext
-            .set(codeLenses && codeLenses.length > 0)
-            .catch((ex) => logger.debug('Failed to set jupyter.HasCodeCells context', ex));
+        const hasCodeCells = codeLenses && codeLenses.length > 0;
+        hasCellsContext.set(hasCodeCells).catch((ex) => logger.debug('Failed to set jupyter.HasCodeCells context', ex));
+        this.updateOwnerContextKey(hasCodeCells);
 
         // Don't provide any code lenses if we have not enabled data science
         const settings = this.configuration.getSettings(document.uri);
diff --git a/src/interactive-window/editor-integration/codelensprovider.unit.test.ts b/src/interactive-window/editor-integration/codelensprovider.unit.test.ts
@@ -1,23 +1,22 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import { anything, when } from 'ts-mockito';
+import { anything, verify, when } from 'ts-mockito';
 import * as TypeMoq from 'typemoq';
-import { CancellationTokenSource, Disposable, EventEmitter, TextDocument, Uri } from 'vscode';
+import { CancellationTokenSource, CodeLens, Disposable, EventEmitter, TextDocument, Uri, Range } from 'vscode';
 
 import { IDebugService } from '../../platform/common/application/types';
 import { IConfigurationService, IWatchableJupyterSettings } from '../../platform/common/types';
 import { DataScienceCodeLensProvider } from '../../interactive-window/editor-integration/codelensprovider';
 import { IServiceContainer } from '../../platform/ioc/types';
-import { ICodeWatcher, IDataScienceCodeLensProvider } from '../../interactive-window/editor-integration/types';
+import { ICodeWatcher } from '../../interactive-window/editor-integration/types';
 import { IDebugLocationTracker } from '../../notebooks/debugger/debuggingTypes';
 import { mockedVSCodeNamespaces } from '../../test/vscode-mock';
 
 // eslint-disable-next-line
 suite('DataScienceCodeLensProvider Unit Tests', () => {
     let serviceContainer: TypeMoq.IMock<IServiceContainer>;
     let configurationService: TypeMoq.IMock<IConfigurationService>;
-    let codeLensProvider: IDataScienceCodeLensProvider;
     let pythonSettings: TypeMoq.IMock<IWatchableJupyterSettings>;
     let debugService: TypeMoq.IMock<IDebugService>;
     let debugLocationTracker: TypeMoq.IMock<IDebugLocationTracker>;
@@ -37,16 +36,17 @@ suite('DataScienceCodeLensProvider Unit Tests', () => {
         configurationService.setup((c) => c.getSettings(TypeMoq.It.isAny())).returns(() => pythonSettings.object);
         when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything(), anything())).thenResolve();
         debugService.setup((d) => d.activeDebugSession).returns(() => undefined);
-        codeLensProvider = new DataScienceCodeLensProvider(
+    });
+
+    function provideCodeLensesForOneDoc(codeLenses: CodeLens[] = []) {
+        const codeLensProvider = new DataScienceCodeLensProvider(
             serviceContainer.object,
             debugLocationTracker.object,
             configurationService.object,
             disposables,
             debugService.object
         );
-    });
 
-    test('Initialize Code Lenses one document', async () => {
         // Create our document
         const document = TypeMoq.Mock.ofType<TextDocument>();
         const uri = Uri.file('test.py');
@@ -57,21 +57,35 @@ suite('DataScienceCodeLensProvider Unit Tests', () => {
         const targetCodeWatcher = TypeMoq.Mock.ofType<ICodeWatcher>();
         targetCodeWatcher
             .setup((tc) => tc.getCodeLenses())
-            .returns(() => [])
+            .returns(() => codeLenses)
             .verifiable(TypeMoq.Times.once());
         serviceContainer
             .setup((c) => c.get(TypeMoq.It.isValue(ICodeWatcher)))
             .returns(() => targetCodeWatcher.object)
             .verifiable(TypeMoq.Times.once());
         when(mockedVSCodeNamespaces.workspace.textDocuments).thenReturn([document.object]);
 
-        await codeLensProvider.provideCodeLenses(document.object, tokenSource.token);
+        codeLensProvider.provideCodeLenses(document.object, tokenSource.token);
+
+        return targetCodeWatcher;
+    }
+
+    test('Initialize Code Lenses one document', async () => {
+        const targetCodeWatcher = provideCodeLensesForOneDoc();
 
         targetCodeWatcher.verifyAll();
         serviceContainer.verifyAll();
     });
 
     test('Initialize Code Lenses same doc called', async () => {
+        const codeLensProvider = new DataScienceCodeLensProvider(
+            serviceContainer.object,
+            debugLocationTracker.object,
+            configurationService.object,
+            disposables,
+            debugService.object
+        );
+
         // Create our document
         const document = TypeMoq.Mock.ofType<TextDocument>();
         const uri = Uri.file('test.py');
@@ -94,15 +108,23 @@ suite('DataScienceCodeLensProvider Unit Tests', () => {
             .verifiable(TypeMoq.Times.once());
         when(mockedVSCodeNamespaces.workspace.textDocuments).thenReturn([document.object]);
 
-        await codeLensProvider.provideCodeLenses(document.object, tokenSource.token);
-        await codeLensProvider.provideCodeLenses(document.object, tokenSource.token);
+        codeLensProvider.provideCodeLenses(document.object, tokenSource.token);
+        codeLensProvider.provideCodeLenses(document.object, tokenSource.token);
 
         // getCodeLenses should be called twice, but getting the code watcher only once due to same doc
         targetCodeWatcher.verifyAll();
         serviceContainer.verifyAll();
     });
 
     test('Initialize Code Lenses different documents', async () => {
+        const codeLensProvider = new DataScienceCodeLensProvider(
+            serviceContainer.object,
+            debugLocationTracker.object,
+            configurationService.object,
+            disposables,
+            debugService.object
+        );
+
         // Create our document
         const uri1 = Uri.file('test.py');
         const document1 = TypeMoq.Mock.ofType<TextDocument>();
@@ -134,13 +156,40 @@ suite('DataScienceCodeLensProvider Unit Tests', () => {
 
         when(mockedVSCodeNamespaces.workspace.textDocuments).thenReturn([document1.object, document2.object]);
 
-        await codeLensProvider.provideCodeLenses(document1.object, tokenSource.token);
-        await codeLensProvider.provideCodeLenses(document1.object, tokenSource.token);
-        await codeLensProvider.provideCodeLenses(document2.object, tokenSource.token);
+        codeLensProvider.provideCodeLenses(document1.object, tokenSource.token);
+        codeLensProvider.provideCodeLenses(document1.object, tokenSource.token);
+        codeLensProvider.provideCodeLenses(document2.object, tokenSource.token);
 
         // service container get should be called three times as the names and versions don't match
         targetCodeWatcher1.verifyAll();
         targetCodeWatcher2.verifyAll();
         serviceContainer.verifyAll();
     });
+
+    test('Having code lenses will update context keys to true', async () => {
+        pythonSettings.setup((p) => p.sendSelectionToInteractiveWindow).returns(() => true);
+
+        provideCodeLensesForOneDoc([new CodeLens({} as Range)]);
+
+        verify(mockedVSCodeNamespaces.commands.executeCommand('setContext', 'jupyter.ownsSelection', true)).atLeast(1);
+        verify(mockedVSCodeNamespaces.commands.executeCommand('setContext', 'jupyter.hascodecells', true)).atLeast(1);
+    });
+
+    test('Having no code lenses will set context keys to false', async () => {
+        pythonSettings.setup((p) => p.sendSelectionToInteractiveWindow).returns(() => false);
+
+        provideCodeLensesForOneDoc([]);
+
+        verify(mockedVSCodeNamespaces.commands.executeCommand('setContext', 'jupyter.ownsSelection', false)).atLeast(1);
+        verify(mockedVSCodeNamespaces.commands.executeCommand('setContext', 'jupyter.hascodecells', false)).atLeast(1);
+    });
+
+    test('Having no code lenses but ownership setting true will set context keys correctly', async () => {
+        pythonSettings.setup((p) => p.sendSelectionToInteractiveWindow).returns(() => true);
+
+        provideCodeLensesForOneDoc([]);
+
+        verify(mockedVSCodeNamespaces.commands.executeCommand('setContext', 'jupyter.ownsSelection', true)).atLeast(1);
+        verify(mockedVSCodeNamespaces.commands.executeCommand('setContext', 'jupyter.hascodecells', false)).atLeast(1);
+    });
 });
diff --git a/src/standalone/activation/globalActivation.ts b/src/standalone/activation/globalActivation.ts
@@ -3,12 +3,7 @@
 
 import { inject, injectable, multiInject, optional } from 'inversify';
 import { ContextKey } from '../../platform/common/contextKey';
-import {
-    IConfigurationService,
-    IDataScienceCommandListener,
-    IDisposable,
-    IDisposableRegistry
-} from '../../platform/common/types';
+import { IDataScienceCommandListener, IDisposable, IDisposableRegistry } from '../../platform/common/types';
 import { noop } from '../../platform/common/utils/misc';
 import { EditorContexts } from '../../platform/common/constants';
 import { IExtensionSyncActivationService } from '../../platform/activation/types';
@@ -25,7 +20,6 @@ export class GlobalActivation implements IExtensionSyncActivationService {
     private startTime: number = Date.now();
     constructor(
         @inject(IDisposableRegistry) private disposableRegistry: IDisposableRegistry,
-        @inject(IConfigurationService) private configuration: IConfigurationService,
         @inject(IRawNotebookSupportedService)
         @optional()
         private rawSupported: IRawNotebookSupportedService | undefined,
@@ -38,9 +32,6 @@ export class GlobalActivation implements IExtensionSyncActivationService {
     }
 
     public activate() {
-        // Set our initial settings and sign up for changes
-        this.onSettingsChanged();
-        this.changeHandler = this.configuration.getSettings(undefined).onDidChange(this.onSettingsChanged.bind(this));
         this.disposableRegistry.push(this);
 
         // Figure out the ZMQ available context key
@@ -58,13 +49,6 @@ export class GlobalActivation implements IExtensionSyncActivationService {
         }
     }
 
-    private onSettingsChanged = () => {
-        const settings = this.configuration.getSettings(undefined);
-        const ownsSelection = settings.sendSelectionToInteractiveWindow;
-        const editorContext = new ContextKey(EditorContexts.OwnsSelection);
-        editorContext.set(ownsSelection).catch(noop);
-    };
-
     private computeZmqAvailable() {
         const zmqContext = new ContextKey(EditorContexts.ZmqAvailable);
         zmqContext.set(this.rawSupported ? this.rawSupported.isSupported : false).then(noop, noop);
diff --git a/src/test/datascience/datascience.unit.test.ts b/src/test/datascience/datascience.unit.test.ts
@@ -8,14 +8,12 @@ import { anything, instance, mock, when } from 'ts-mockito';
 import { JupyterSettings } from '../../platform/common/configSettings';
 import { ConfigurationService } from '../../platform/common/configuration/service.node';
 import { IConfigurationService, IWatchableJupyterSettings } from '../../platform/common/types';
-import { GlobalActivation } from '../../standalone/activation/globalActivation';
 import { RawNotebookSupportedService } from '../../kernels/raw/session/rawNotebookSupportedService.node';
 import { IRawNotebookSupportedService } from '../../kernels/raw/types';
 import { pruneCell } from '../../platform/common/utils';
 
 /* eslint-disable  */
 suite('Tests', () => {
-    let dataScience: GlobalActivation;
     let configService: IConfigurationService;
     let settings: IWatchableJupyterSettings;
     let onDidChangeSettings: sinon.SinonStub;
@@ -25,30 +23,12 @@ suite('Tests', () => {
         settings = mock(JupyterSettings);
         rawNotebookSupported = mock(RawNotebookSupportedService);
 
-        dataScience = new GlobalActivation(
-            // eslint-disable-next-line @typescript-eslint/no-explicit-any
-            [] as any,
-            instance(configService),
-            instance(rawNotebookSupported),
-            [] as any
-        );
-
         onDidChangeSettings = sinon.stub();
         when(configService.getSettings(anything())).thenReturn(instance(settings));
         when(settings.onDidChange).thenReturn(onDidChangeSettings);
         when(rawNotebookSupported.isSupported).thenReturn(true);
     });
 
-    suite('Activate', () => {
-        setup(async () => {
-            await dataScience.activate();
-        });
-
-        test('Should add handler for Settings Changed', async () => {
-            assert.ok(onDidChangeSettings.calledOnce);
-        });
-    });
-
     suite('Cell pruning', () => {
         test('Remove output and execution count from non code', () => {
             const cell: nbformat.ICell = {
