Skip to content

Commit a205041

Browse files
Copilotluabud
Copilot
and
luabud
committed
Implement DataFrame tracker for Jupyter extension suggestion
Co-authored-by: luabud <[email protected]>
1 parent c2db48c commit a205041

File tree

5 files changed

+270
-1
lines changed

5 files changed

+270
-1
lines changed

‎src/client/debugger/extension/adapter/activator.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import { IConfigurationService, IDisposableRegistry } from '../../../common/type
1010
import { ICommandManager } from '../../../common/application/types';
1111
import { DebuggerTypeName } from '../../constants';
1212
import { IAttachProcessProviderFactory } from '../attachQuickPick/types';
13-
import { IDebugAdapterDescriptorFactory, IDebugSessionLoggingFactory, IOutdatedDebuggerPromptFactory } from '../types';
13+
import { IDebugAdapterDescriptorFactory, IDebugSessionLoggingFactory, IOutdatedDebuggerPromptFactory, IDataFrameTrackerFactory } from '../types';
1414

1515
@injectable()
1616
export class DebugAdapterActivator implements IExtensionSingleActivationService {
@@ -22,6 +22,7 @@ export class DebugAdapterActivator implements IExtensionSingleActivationService
2222
@inject(IDebugAdapterDescriptorFactory) private descriptorFactory: IDebugAdapterDescriptorFactory,
2323
@inject(IDebugSessionLoggingFactory) private debugSessionLoggingFactory: IDebugSessionLoggingFactory,
2424
@inject(IOutdatedDebuggerPromptFactory) private debuggerPromptFactory: IOutdatedDebuggerPromptFactory,
25+
@inject(IDataFrameTrackerFactory) private dataFrameTrackerFactory: IDataFrameTrackerFactory,
2526
@inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
2627
@inject(IAttachProcessProviderFactory)
2728
private readonly attachProcessProviderFactory: IAttachProcessProviderFactory,
@@ -35,6 +36,9 @@ export class DebugAdapterActivator implements IExtensionSingleActivationService
3536
this.disposables.push(
3637
this.debugService.registerDebugAdapterTrackerFactory(DebuggerTypeName, this.debuggerPromptFactory),
3738
);
39+
this.disposables.push(
40+
this.debugService.registerDebugAdapterTrackerFactory(DebuggerTypeName, this.dataFrameTrackerFactory),
41+
);
3842

3943
this.disposables.push(
4044
this.debugService.registerDebugAdapterDescriptorFactory(DebuggerTypeName, this.descriptorFactory),

‎src/client/debugger/extension/adapter/dataFrameTracker.ts‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
'use strict';
5+
6+
import { inject, injectable } from 'inversify';
7+
import {
8+
DebugAdapterTracker,
9+
DebugAdapterTrackerFactory,
10+
DebugSession,
11+
ProviderResult,
12+
window,
13+
l10n,
14+
commands,
15+
} from 'vscode';
16+
import { DebugProtocol } from 'vscode-debugprotocol';
17+
18+
import { IExtensions } from '../../../common/types';
19+
import { JUPYTER_EXTENSION_ID } from '../../../common/constants';
20+
21+
/**
22+
* Debug adapter tracker that monitors for dataframe-like variables
23+
* and suggests installing the Jupyter extension when they are detected
24+
* but the Jupyter extension is not installed.
25+
*/
26+
class DataFrameVariableTracker implements DebugAdapterTracker {
27+
private readonly extensions: IExtensions;
28+
private hasNotifiedAboutJupyter = false;
29+
30+
// Types that are considered dataframe-like
31+
private readonly dataFrameTypes = [
32+
'pandas.core.frame.DataFrame',
33+
'pandas.DataFrame',
34+
'polars.DataFrame',
35+
'cudf.DataFrame',
36+
'dask.dataframe.core.DataFrame',
37+
'modin.pandas.DataFrame',
38+
'vaex.dataframe.DataFrame',
39+
'geopandas.geodataframe.GeoDataFrame',
40+
];
41+
42+
constructor(_session: DebugSession, extensions: IExtensions) {
43+
this.extensions = extensions;
44+
}
45+
46+
public onDidSendMessage(message: DebugProtocol.Message): void {
47+
if (this.hasNotifiedAboutJupyter) {
48+
return; // Only notify once per debug session
49+
}
50+
51+
// Check if this is a variables response
52+
if ('type' in message && message.type === 'response' && 'command' in message && message.command === 'variables') {
53+
const response = message as unknown as DebugProtocol.VariablesResponse;
54+
if (response.success && response.body?.variables) {
55+
this.checkForDataFrameVariables(response.body.variables);
56+
}
57+
}
58+
}
59+
60+
private checkForDataFrameVariables(variables: DebugProtocol.Variable[]): boolean {
61+
// Check if any variable is a dataframe-like object
62+
const hasDataFrame = variables.some((variable) =>
63+
this.dataFrameTypes.some((dfType) =>
64+
variable.type?.includes(dfType) || variable.value?.includes(dfType)
65+
)
66+
);
67+
68+
if (hasDataFrame) {
69+
this.checkAndNotifyJupyterExtension();
70+
}
71+
72+
return hasDataFrame;
73+
}
74+
75+
private checkAndNotifyJupyterExtension(): void {
76+
// Check if Jupyter extension is installed
77+
const jupyterExtension = this.extensions.getExtension(JUPYTER_EXTENSION_ID);
78+
79+
if (!jupyterExtension) {
80+
this.hasNotifiedAboutJupyter = true;
81+
this.showJupyterInstallNotification();
82+
}
83+
}
84+
85+
private showJupyterInstallNotification(): void {
86+
const message = l10n.t('Install Jupyter extension to inspect dataframe objects in the data viewer.');
87+
const installAction = l10n.t('Install Jupyter Extension');
88+
const dismissAction = l10n.t('Dismiss');
89+
90+
window.showInformationMessage(message, installAction, dismissAction).then((selection) => {
91+
if (selection === installAction) {
92+
// Open the extension marketplace for the Jupyter extension
93+
commands.executeCommand('extension.open', JUPYTER_EXTENSION_ID);
94+
}
95+
});
96+
}
97+
}
98+
99+
@injectable()
100+
export class DataFrameTrackerFactory implements DebugAdapterTrackerFactory {
101+
constructor(@inject(IExtensions) private readonly extensions: IExtensions) {}
102+
103+
public createDebugAdapterTracker(session: DebugSession): ProviderResult<DebugAdapterTracker> {
104+
return new DataFrameVariableTracker(session, this.extensions);
105+
}
106+
}

‎src/client/debugger/extension/serviceRegistry.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import { DebugAdapterActivator } from './adapter/activator';
1010
import { DebugAdapterDescriptorFactory } from './adapter/factory';
1111
import { DebugSessionLoggingFactory } from './adapter/logging';
1212
import { OutdatedDebuggerPromptFactory } from './adapter/outdatedDebuggerPrompt';
13+
import { DataFrameTrackerFactory } from './adapter/dataFrameTracker';
1314
import { AttachProcessProviderFactory } from './attachQuickPick/factory';
1415
import { IAttachProcessProviderFactory } from './attachQuickPick/types';
1516
import { PythonDebugConfigurationService } from './configuration/debugConfigurationService';
@@ -26,6 +27,7 @@ import {
2627
IDebugConfigurationService,
2728
IDebugSessionLoggingFactory,
2829
IOutdatedDebuggerPromptFactory,
30+
IDataFrameTrackerFactory,
2931
} from './types';
3032

3133
export function registerTypes(serviceManager: IServiceManager): void {
@@ -62,6 +64,10 @@ export function registerTypes(serviceManager: IServiceManager): void {
6264
IOutdatedDebuggerPromptFactory,
6365
OutdatedDebuggerPromptFactory,
6466
);
67+
serviceManager.addSingleton<IDataFrameTrackerFactory>(
68+
IDataFrameTrackerFactory,
69+
DataFrameTrackerFactory,
70+
);
6571
serviceManager.addSingleton<IAttachProcessProviderFactory>(
6672
IAttachProcessProviderFactory,
6773
AttachProcessProviderFactory,

‎src/client/debugger/extension/types.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,10 @@ export const IOutdatedDebuggerPromptFactory = Symbol('IOutdatedDebuggerPromptFac
1919

2020
export interface IOutdatedDebuggerPromptFactory extends DebugAdapterTrackerFactory {}
2121

22+
export const IDataFrameTrackerFactory = Symbol('IDataFrameTrackerFactory');
23+
24+
export interface IDataFrameTrackerFactory extends DebugAdapterTrackerFactory {}
25+
2226
export enum PythonPathSource {
2327
launchJson = 'launch.json',
2428
settingsJson = 'settings.json',

‎src/test/debugger/extension/adapter/dataFrameTracker.unit.test.ts‎

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
'use strict';
5+
6+
import { expect } from 'chai';
7+
import { instance, mock, verify, when } from 'ts-mockito';
8+
import { Extension } from 'vscode';
9+
import { DebugProtocol } from 'vscode-debugprotocol';
10+
11+
import { IExtensions } from '../../../../client/common/types';
12+
import { JUPYTER_EXTENSION_ID } from '../../../../client/common/constants';
13+
import { DataFrameTrackerFactory } from '../../../../client/debugger/extension/adapter/dataFrameTracker';
14+
15+
suite('DataFrame Tracker', () => {
16+
let extensions: IExtensions;
17+
let mockExtensions: IExtensions;
18+
let trackerFactory: DataFrameTrackerFactory;
19+
20+
setup(() => {
21+
mockExtensions = mock<IExtensions>();
22+
extensions = instance(mockExtensions);
23+
trackerFactory = new DataFrameTrackerFactory(extensions);
24+
});
25+
26+
test('Should create debug adapter tracker', () => {
27+
const mockSession = {} as any;
28+
const tracker = trackerFactory.createDebugAdapterTracker(mockSession);
29+
expect(tracker).to.not.be.undefined;
30+
});
31+
32+
test('Should detect pandas DataFrame variable', () => {
33+
const mockSession = {} as any;
34+
const tracker = trackerFactory.createDebugAdapterTracker(mockSession) as any;
35+
36+
// Mock Jupyter extension not being installed
37+
when(mockExtensions.getExtension(JUPYTER_EXTENSION_ID)).thenReturn(undefined);
38+
39+
const variablesResponse: DebugProtocol.VariablesResponse = {
40+
type: 'response',
41+
seq: 1,
42+
request_seq: 1,
43+
success: true,
44+
command: 'variables',
45+
body: {
46+
variables: [
47+
{
48+
name: 'df',
49+
value: '<pandas.core.frame.DataFrame object>',
50+
type: 'pandas.core.frame.DataFrame',
51+
variablesReference: 0,
52+
},
53+
{
54+
name: 'x',
55+
value: '42',
56+
type: 'int',
57+
variablesReference: 0,
58+
}
59+
]
60+
}
61+
};
62+
63+
// This should detect the DataFrame and try to show notification
64+
tracker.onDidSendMessage(variablesResponse);
65+
66+
// Verify that extension check was called
67+
verify(mockExtensions.getExtension(JUPYTER_EXTENSION_ID)).once();
68+
});
69+
70+
test('Should not show notification if Jupyter extension is installed', () => {
71+
const mockSession = {} as any;
72+
const tracker = trackerFactory.createDebugAdapterTracker(mockSession) as any;
73+
74+
// Mock Jupyter extension being installed
75+
const mockJupyterExt = mock<Extension<any>>();
76+
when(mockExtensions.getExtension(JUPYTER_EXTENSION_ID)).thenReturn(instance(mockJupyterExt));
77+
78+
const variablesResponse: DebugProtocol.VariablesResponse = {
79+
type: 'response',
80+
seq: 1,
81+
request_seq: 1,
82+
success: true,
83+
command: 'variables',
84+
body: {
85+
variables: [
86+
{
87+
name: 'df',
88+
value: '<pandas.core.frame.DataFrame object>',
89+
type: 'pandas.core.frame.DataFrame',
90+
variablesReference: 0,
91+
}
92+
]
93+
}
94+
};
95+
96+
tracker.onDidSendMessage(variablesResponse);
97+
98+
// Verify that extension check was called but no notification should show
99+
verify(mockExtensions.getExtension(JUPYTER_EXTENSION_ID)).once();
100+
});
101+
102+
test('Should detect various dataframe types', () => {
103+
// This test verifies that the dataFrameTypes array contains the expected types
104+
const expectedTypes = [
105+
'pandas.core.frame.DataFrame',
106+
'pandas.DataFrame',
107+
'polars.DataFrame',
108+
'cudf.DataFrame',
109+
'dask.dataframe.core.DataFrame',
110+
'modin.pandas.DataFrame',
111+
'vaex.dataframe.DataFrame',
112+
'geopandas.geodataframe.GeoDataFrame',
113+
];
114+
115+
expectedTypes.forEach(expectedType => {
116+
// Verify each expected type would be matched by at least one pattern
117+
const hasMatch = expectedTypes.some(pattern => expectedType.includes(pattern));
118+
expect(hasMatch).to.be.true;
119+
});
120+
});
121+
122+
test('Should not detect non-dataframe variables', () => {
123+
const nonDataFrameTypes = [
124+
'str',
125+
'int',
126+
'list',
127+
'dict',
128+
'numpy.ndarray',
129+
'matplotlib.figure.Figure',
130+
];
131+
132+
const dataFrameTypes = [
133+
'pandas.core.frame.DataFrame',
134+
'pandas.DataFrame',
135+
'polars.DataFrame',
136+
'cudf.DataFrame',
137+
'dask.dataframe.core.DataFrame',
138+
'modin.pandas.DataFrame',
139+
'vaex.dataframe.DataFrame',
140+
'geopandas.geodataframe.GeoDataFrame',
141+
];
142+
143+
nonDataFrameTypes.forEach(varType => {
144+
// These should not be detected as dataframes
145+
const hasMatch = dataFrameTypes.some(dfType => varType.includes(dfType));
146+
expect(hasMatch).to.be.false;
147+
});
148+
});
149+
});

0 commit comments

Comments
 (0)