Revert settings changes when opting out of Native NB Experiment

src/client/datascience/notebook/integration.ts

@@ -12,14 +12,16 @@ import {
 } from '../../common/application/types';
 import { NotebookEditorSupport } from '../../common/experiments/groups';
 import { traceError } from '../../common/logger';
-import { IDisposableRegistry, IExperimentsManager } from '../../common/types';
+import { IDisposableRegistry, IExperimentsManager, IExtensionContext } from '../../common/types';
 import { DataScience } from '../../common/utils/localize';
 import { noop } from '../../common/utils/misc';
 import { JupyterNotebookView } from './constants';
 import { NotebookKernel } from './notebookKernel';
 import { NotebookOutputRenderer } from './renderer';
 import { INotebookContentProvider } from './types';
 
+const EditorAssociationUpdatedKey = 'EditorAssociationUpdatedToUseNotebooks';
+
 /**
  * This class basically registers the necessary providers and the like with VSC.
  * I.e. this is where we integrate our stuff with VS Code via their extension endpoints.
@@ -36,14 +38,19 @@ export class NotebookIntegration implements IExtensionSingleActivationService {
         @inject(NotebookOutputRenderer) private readonly renderer: NotebookOutputRenderer,
         @inject(IApplicationEnvironment) private readonly env: IApplicationEnvironment,
         @inject(IApplicationShell) private readonly shell: IApplicationShell,
-        @inject(IWorkspaceService) private readonly workspace: IWorkspaceService
+        @inject(IWorkspaceService) private readonly workspace: IWorkspaceService,
+        @inject(IExtensionContext) private readonly extensionContext: IExtensionContext
     ) {}
     public async activate(): Promise<void> {
         // This condition is temporary.
         // If user belongs to the experiment, then make the necessary changes to package.json.
         // Once the API is final, we won't need to modify the package.json.
         if (this.experiment.inExperiment(NotebookEditorSupport.nativeNotebookExperiment)) {
             await this.enableNotebooks();
+        } else {
+            // Possible user was in experiment, then they opted out. In this case we need to revert the changes made to the settings file.
+            // Again, this is temporary code.
+            await this.disableNotebooks();
         }
         if (this.env.channel !== 'insiders') {
             return;
@@ -95,23 +102,55 @@ export class NotebookIntegration implements IExtensionSingleActivationService {
             return;
         }
 
+        await this.enableDisableEditorAssociation(true);
+    }
+    private async enableDisableEditorAssociation(enable: boolean) {
         // This code is temporary.
         const settings = this.workspace.getConfiguration('workbench', undefined);
         const editorAssociations = settings.get('editorAssociations') as {
             viewType: string;
             filenamePattern: string;
         }[];
 
+        // Update the settings.
         if (
-            !Array.isArray(editorAssociations) ||
-            editorAssociations.length === 0 ||
-            !editorAssociations.find((item) => item.viewType === JupyterNotebookView)
+            enable &&
+            (!Array.isArray(editorAssociations) ||
+                editorAssociations.length === 0 ||
+                !editorAssociations.find((item) => item.viewType === JupyterNotebookView))
         ) {
             editorAssociations.push({
                 viewType: 'jupyter-notebook',
                 filenamePattern: '*.ipynb'
             });
-            await settings.update('editorAssociations', editorAssociations, ConfigurationTarget.Global);
+            await Promise.all([
+                this.extensionContext.globalState.update(EditorAssociationUpdatedKey, true),
+                settings.update('editorAssociations', editorAssociations, ConfigurationTarget.Global)
+            ]);
+        }
+
+        // Revert the settings.
+        if (
+            !enable &&
+            this.extensionContext.globalState.get<boolean>(EditorAssociationUpdatedKey, false) &&
+            Array.isArray(editorAssociations) &&
+            editorAssociations.find((item) => item.viewType === JupyterNotebookView)
+        ) {
+            const updatedSettings = editorAssociations.filter((item) => item.viewType !== JupyterNotebookView);
+            await Promise.all([
+                this.extensionContext.globalState.update(EditorAssociationUpdatedKey, false),
+                settings.update('editorAssociations', updatedSettings, ConfigurationTarget.Global)
+            ]);
+        }
+    }
+    private async disableNotebooks() {
+        if (this.env.channel === 'stable') {
+            return;
+        }
+        // If we never modified the settings, then nothing to do.
+        if (!this.extensionContext.globalState.get<boolean>(EditorAssociationUpdatedKey, false)) {
+            return;
         }
+        await this.enableDisableEditorAssociation(false);
     }
 }