Add code folding regions for # %% cells within python files

news/1 Enhancements/1527.md

@@ -0,0 +1 @@
+Add code folding regions for `# %%` cells within python files.

src/interactive-window/editor-integration/pythonCellFoldingProvider.ts

@@ -0,0 +1,38 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { inject, injectable } from 'inversify';
+import {
+    TextDocument,
+    FoldingContext,
+    CancellationToken,
+    ProviderResult,
+    FoldingRange,
+    FoldingRangeKind
+} from 'vscode';
+import { IDataScienceCodeLensProvider, IPythonCellFoldingProvider } from './types';
+
+@injectable()
+export class PythonCellFoldingProvider implements IPythonCellFoldingProvider {
+    constructor(
+        @inject(IDataScienceCodeLensProvider) private dataScienceCodeLensProvider: IDataScienceCodeLensProvider
+    ) {}
+
+    provideFoldingRanges(
+        document: TextDocument,
+        _context: FoldingContext,
+        token: CancellationToken
+    ): ProviderResult<FoldingRange[]> {
+        const codeWatcher = this.dataScienceCodeLensProvider.getCodeWatcher(document);
+        if (codeWatcher) {
+            const codeLenses = codeWatcher.getCodeLenses();
+            if (token.isCancellationRequested) {
+                return undefined;
+            }
+            return codeLenses.map((codeLens) => {
+                return new FoldingRange(codeLens.range.start.line, codeLens.range.end.line, FoldingRangeKind.Region);
+            });
+        }
+        return undefined;
+    }
+}

src/interactive-window/editor-integration/types.ts

@@ -3,7 +3,17 @@
 
 'use strict';
 
-import { Event, CodeLens, CodeLensProvider, Uri, TextEditor, Range, TextDocument, NotebookDocument } from 'vscode';
+import {
+    Event,
+    CodeLens,
+    CodeLensProvider,
+    Uri,
+    TextEditor,
+    Range,
+    TextDocument,
+    NotebookDocument,
+    FoldingRangeProvider
+} from 'vscode';
 import { ICellRange, IDisposable } from '../../platform/common/types';
 
 // Wraps the vscode CodeLensProvider base class
@@ -12,6 +22,9 @@ export interface IDataScienceCodeLensProvider extends CodeLensProvider {
     getCodeWatcher(document: TextDocument): ICodeWatcher | undefined;
 }
 
+export const IPythonCellFoldingProvider = Symbol('IPythonCellFoldingProvider');
+export interface IPythonCellFoldingProvider extends FoldingRangeProvider {}
+
 // Wraps the Code Watcher API
 export const ICodeWatcher = Symbol('ICodeWatcher');
 export interface ICodeWatcher extends IDisposable {

src/interactive-window/serviceRegistry.node.ts

@@ -21,7 +21,8 @@ import {
     ICodeLensFactory,
     IDataScienceCodeLensProvider,
     IGeneratedCodeStorageFactory,
-    ICodeGeneratorFactory
+    ICodeGeneratorFactory,
+    IPythonCellFoldingProvider
 } from './editor-integration/types';
 import { GeneratedCodeStorageManager } from './generatedCodeStoreManager';
 import { InteractiveWindowTracebackFormatter } from './outputs/tracebackFormatter';
@@ -30,6 +31,7 @@ import { InteractiveWindowDebugger } from './debugger/interactiveWindowDebugger.
 import { InteractiveWindowDebuggingManager } from './debugger/jupyter/debuggingManager';
 import { BANNER_NAME_INTERACTIVE_SHIFTENTER, InteractiveShiftEnterBanner } from './shiftEnterBanner';
 import { InteractiveWindowDebuggingStartupCodeProvider } from './debugger/startupCodeProvider';
+import { PythonCellFoldingProvider } from './editor-integration/pythonCellFoldingProvider';
 
 export function registerTypes(serviceManager: IServiceManager) {
     serviceManager.addSingleton<IInteractiveWindowProvider>(IInteractiveWindowProvider, InteractiveWindowProvider);
@@ -41,6 +43,7 @@ export function registerTypes(serviceManager: IServiceManager) {
         IDataScienceCodeLensProvider,
         DataScienceCodeLensProvider
     );
+    serviceManager.addSingleton<IPythonCellFoldingProvider>(IPythonCellFoldingProvider, PythonCellFoldingProvider);
     serviceManager.addSingleton<IExtensionSingleActivationService>(IExtensionSingleActivationService, Decorator);
     serviceManager.addSingleton<IExtensionSyncActivationService>(
         IExtensionSyncActivationService,

src/standalone/activation/globalActivation.ts

@@ -20,7 +20,10 @@ import { debounceAsync, swallowExceptions } from '../../platform/common/utils/de
 import { noop } from '../../platform/common/utils/misc';
 import { EditorContexts } from '../../platform/common/constants';
 import { IExtensionSingleActivationService } from '../../platform/activation/types';
-import { IDataScienceCodeLensProvider } from '../../interactive-window/editor-integration/types';
+import {
+    IDataScienceCodeLensProvider,
+    IPythonCellFoldingProvider
+} from '../../interactive-window/editor-integration/types';
 import { IRawNotebookSupportedService } from '../../kernels/raw/types';
 import { hasCells } from '../../interactive-window/editor-integration/cellFactory';
 import { sendTelemetryEvent } from '../../telemetry';
@@ -41,6 +44,9 @@ export class GlobalActivation implements IExtensionSingleActivationService {
         @inject(IDataScienceCodeLensProvider)
         @optional()
         private dataScienceCodeLensProvider: IDataScienceCodeLensProvider | undefined,
+        @optional()
+        @inject(IPythonCellFoldingProvider)
+        private pythonCellFoldingProvider: IPythonCellFoldingProvider | undefined,
         @inject(IConfigurationService) private configuration: IConfigurationService,
         @inject(IDocumentManager) private documentManager: IDocumentManager,
         @inject(IWorkspaceService) private workspace: IWorkspaceService,
@@ -61,6 +67,11 @@ export class GlobalActivation implements IExtensionSingleActivationService {
                 vscode.languages.registerCodeLensProvider([PYTHON_FILE_ANY_SCHEME], this.dataScienceCodeLensProvider)
             );
         }
+        if (this.pythonCellFoldingProvider) {
+            this.extensionContext.subscriptions.push(
+                vscode.languages.registerFoldingRangeProvider([PYTHON_FILE_ANY_SCHEME], this.pythonCellFoldingProvider)
+            );
+        }
 
         // Set our initial settings and sign up for changes
         this.onSettingsChanged();

src/test/datascience/datascience.unit.test.ts

@@ -49,6 +49,7 @@ suite('DataScience Tests', () => {
             // eslint-disable-next-line @typescript-eslint/no-explicit-any
             { subscriptions: [] } as any,
             instance(dataScienceCodeLensProvider),
+            undefined,
             instance(configService),
             instance(docManager),
             instance(workspaceService),