Error line numbers coming out with invalid values (#11234)

* Working inside cell hash provider

* Add more unit tests

* Add news entry

* Fix sonar warnings

Author: rchiodo

news/2 Fixes/10708.md

@@ -0,0 +1 @@
+Make error links in exception tracebacks support multiple cells in the stack and extra spaces.

src/client/datascience/editor-integration/cellhashprovider.ts

@@ -1,8 +1,10 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 'use strict';
+import type { KernelMessage } from '@jupyterlab/services';
 import * as hashjs from 'hash.js';
 import { inject, injectable, multiInject, optional } from 'inversify';
+import stripAnsi from 'strip-ansi';
 import { Event, EventEmitter, Position, Range, TextDocumentChangeEvent, TextDocumentContentChangeEvent } from 'vscode';
 
 import { splitMultilineString } from '../../../datascience-ui/common';
@@ -23,12 +25,18 @@ import {
     INotebookExecutionLogger
 } from '../types';
 
+// tslint:disable-next-line:no-require-imports no-var-requires
+const _escapeRegExp = require('lodash/escapeRegExp') as typeof import('lodash/escapeRegExp'); // NOSONAR
+const LineNumberMatchRegex = /(;32m[ ->]*?)(\d+)/g;
+
 interface IRangedCellHash extends ICellHash {
     code: string;
     startOffset: number;
     endOffset: number;
     deleted: boolean;
     realCode: string;
+    trimmedRightCode: string;
+    firstNonBlankLineIndex: number; // zero based. First non blank line of the real code.
 }
 
 // This class provides hashes for debugging jupyter cells. Call getHashes just before starting debugging to compute all of the
@@ -45,6 +53,7 @@ export class CellHashProvider implements ICellHashProvider, INotebookExecutionLo
     private executionCount: number = 0;
     private hashes: Map<string, IRangedCellHash[]> = new Map<string, IRangedCellHash[]>();
     private updateEventEmitter: EventEmitter<void> = new EventEmitter<void>();
+    private traceBackRegexes = new Map<string, RegExp>();
 
     constructor(
         @inject(IDocumentManager) private documentManager: IDocumentManager,
@@ -59,6 +68,7 @@ export class CellHashProvider implements ICellHashProvider, INotebookExecutionLo
 
     public dispose() {
         this.hashes.clear();
+        this.traceBackRegexes.clear();
     }
 
     public get updated(): Event<void> {
@@ -83,6 +93,7 @@ export class CellHashProvider implements ICellHashProvider, INotebookExecutionLo
 
     public onKernelRestarted() {
         this.hashes.clear();
+        this.traceBackRegexes.clear();
         this.executionCount = 0;
         this.updateEventEmitter.fire();
     }
@@ -112,6 +123,21 @@ export class CellHashProvider implements ICellHashProvider, INotebookExecutionLo
         noop();
     }
 
+    public preHandleIOPub(msg: KernelMessage.IIOPubMessage): KernelMessage.IIOPubMessage {
+        // When an error message comes, rewrite the traceback so we can jump back to the correct
+        // cell. For now this only works with the interactive window
+        if (msg.header.msg_type === 'error') {
+            return {
+                ...msg,
+                content: {
+                    ...msg.content,
+                    traceback: this.modifyTraceback(msg as KernelMessage.IErrorMsg) // NOSONAR
+                }
+            };
+        }
+        return msg;
+    }
+
     public extractExecutableLines(cell: ICell): string[] {
         const cellMatcher = new CellMatcher(this.configService.getSettings(getCellResource(cell)).datascience);
         const lines = splitMultilineString(cell.data.source);
@@ -144,6 +170,12 @@ export class CellHashProvider implements ICellHashProvider, INotebookExecutionLo
             const line = doc.lineAt(trueStartLine);
             const endLine = doc.lineAt(Math.min(trueStartLine + stripped.length - 1, doc.lineCount - 1));
 
+            // Find the first non blank line
+            let firstNonBlankIndex = 0;
+            while (firstNonBlankIndex < stripped.length && stripped[firstNonBlankIndex].trim().length === 0) {
+                firstNonBlankIndex += 1;
+            }
+
             // Use the original values however to track edits. This is what we need
             // to move around
             const startOffset = doc.offsetAt(new Position(cell.line, 0));
@@ -180,11 +212,13 @@ export class CellHashProvider implements ICellHashProvider, INotebookExecutionLo
                 hash: hashjs.sha1().update(hashedCode).digest('hex').substr(0, 12),
                 line: line.lineNumber + 1,
                 endLine: endLine.lineNumber + 1,
+                firstNonBlankLineIndex: firstNonBlankIndex + trueStartLine,
                 executionCount: expectedCount,
                 startOffset,
                 endOffset,
                 deleted: false,
                 code: hashedCode,
+                trimmedRightCode: stripped.map((s) => s.replace(/[ \t\r]+\n$/g, '\n')).join(''),
                 realCode,
                 runtimeLine,
                 id: cell.id
@@ -217,6 +251,15 @@ export class CellHashProvider implements ICellHashProvider, INotebookExecutionLo
             }
             this.hashes.set(cell.file, list);
 
+            // Save a regex to find this file later when looking for
+            // exceptions in output
+            if (!this.traceBackRegexes.has(cell.file)) {
+                const fileDisplayName = this.fileSystem.getDisplayName(cell.file);
+                const escaped = _escapeRegExp(fileDisplayName);
+                const fileMatchRegex = new RegExp(`\\[.*?;32m${escaped}`);
+                this.traceBackRegexes.set(cell.file, fileMatchRegex);
+            }
+
             // Tell listeners we have new hashes.
             if (this.listeners) {
                 const hashes = this.getHashes();
@@ -306,4 +349,60 @@ export class CellHashProvider implements ICellHashProvider, INotebookExecutionLo
         // No breakpoint necessary, start on the first line
         return 1;
     }
+
+    // This function will modify a traceback from an error message.
+    // Tracebacks take a form like so:
+    // "---------------------------------------------------------------------------"
+    // "ZeroDivisionError                         Traceback (most recent call last)"
+    // "d:\Training\SnakePython\foo.py in <module>\n      1 print('some more')\n    ----> 2 cause_error()\n    "
+    // "d:\Training\SnakePython\foo.py in cause_error()\n      3     print('error')\n          4     print('now')\n    ----> 5     print( 1 / 0)\n    "
+    // "ZeroDivisionError: division by zero"
+    // Each item in the array being a stack frame.
+    private modifyTraceback(msg: KernelMessage.IErrorMsg): string[] {
+        // Do one frame at a time.
+        return msg.content.traceback ? msg.content.traceback.map(this.modifyTracebackFrame.bind(this)) : [];
+    }
+
+    private findCellOffset(hashes: IRangedCellHash[] | undefined, codeLines: string): number | undefined {
+        if (hashes) {
+            // Go through all cell code looking for these code lines exactly
+            // (although with right side trimmed as that's what a stack trace does)
+            for (const hash of hashes) {
+                const index = hash.trimmedRightCode.indexOf(codeLines);
+                if (index >= 0) {
+                    // Jupyter isn't counting blank lines at the top so use our
+                    // first non blank line
+                    return hash.firstNonBlankLineIndex;
+                }
+            }
+        }
+        // No hash found
+        return undefined;
+    }
+
+    private modifyTracebackFrame(traceFrame: string): string {
+        // See if this item matches any of our cell files
+        const regexes = [...this.traceBackRegexes.entries()];
+        const match = regexes.find((e) => e[1].test(traceFrame));
+        if (match) {
+            // We have a match, pull out the source lines
+            let sourceLines = '';
+            const regex = /(;32m[ ->]*?)(\d+)(.*)/g;
+            for (let l = regex.exec(traceFrame); l && l.length > 3; l = regex.exec(traceFrame)) {
+                const newLine = stripAnsi(l[3]).substr(1); // Seem to have a space on the front
+                sourceLines = `${sourceLines}${newLine}\n`;
+            }
+
+            // Now attempt to find a cell that matches these source lines
+            const offset = this.findCellOffset(this.hashes.get(match[0]), sourceLines);
+            if (offset !== undefined) {
+                return traceFrame.replace(LineNumberMatchRegex, (_s, prefix, num) => {
+                    const n = parseInt(num, 10);
+                    const newLine = offset + n - 1;
+                    return `${prefix}<a href='file://${match[0]}?line=${newLine}'>${newLine + 1}</a>`;
+                });
+            }
+        }
+        return traceFrame;
+    }
 }

src/client/datascience/jupyter/jupyterNotebook.ts

@@ -37,7 +37,7 @@ import {
     InterruptResult,
     KernelSocketInformation
 } from '../types';
-import { expandWorkingDir, modifyTraceback } from './jupyterUtils';
+import { expandWorkingDir } from './jupyterUtils';
 import { LiveKernelModel } from './kernels/types';
 
 // tslint:disable-next-line: no-require-imports
@@ -940,6 +940,9 @@ export class JupyterNotebookBase implements INotebook {
         // tslint:disable-next-line: no-any
         let result: Promise<any> = Promise.resolve();
 
+        // Let our loggers get a first crack at the message. They may change it
+        this.getLoggers().forEach((f) => (msg = f.preHandleIOPub ? f.preHandleIOPub(msg) : msg));
+
         // tslint:disable-next-line:no-require-imports
         const jupyterLab = require('@jupyterlab/services') as typeof import('@jupyterlab/services');
 
@@ -1366,7 +1369,7 @@ export class JupyterNotebookBase implements INotebook {
             output_type: 'error',
             ename: msg.content.ename,
             evalue: msg.content.evalue,
-            traceback: modifyTraceback(cell.file, this.fs.getDisplayName(cell.file), cell.line, msg.content.traceback)
+            traceback: msg.content.traceback
         };
         this.addToCellData(cell, output, clearState);
         cell.state = CellState.error;

src/client/datascience/jupyter/jupyterUtils.ts

@@ -10,13 +10,9 @@ import { IWorkspaceService } from '../../common/application/types';
 import { IDataScienceSettings } from '../../common/types';
 import { noop } from '../../common/utils/misc';
 import { SystemVariables } from '../../common/variables/systemVariables';
-import { Identifiers } from '../constants';
 import { getJupyterConnectionDisplayName } from '../jupyter/jupyterConnection';
 import { IConnection } from '../types';
 
-// tslint:disable-next-line:no-require-imports no-var-requires
-const _escapeRegExp = require('lodash/escapeRegExp') as typeof import('lodash/escapeRegExp');
-
 export function expandWorkingDir(
     workingDir: string | undefined,
     launchingFile: string,
@@ -62,44 +58,3 @@ export function createRemoteConnectionInfo(uri: string, settings: IDataScienceSe
         dispose: noop
     };
 }
-
-const LineMatchRegex = /(;32m[ ->]*?)(\d+)/g;
-const IPythonMatchRegex = /(<ipython-input.*?>)/g;
-
-function modifyLineNumbers(entry: string, file: string, startLine: number): string {
-    return entry.replace(LineMatchRegex, (_s, prefix, num) => {
-        const n = parseInt(num, 10);
-        const newLine = startLine + n;
-        return `${prefix}<a href='file://${file}?line=${newLine}'>${newLine + 1}</a>`;
-    });
-}
-
-function modifyTracebackEntry(
-    fileMatchRegex: RegExp,
-    file: string,
-    fileDisplayName: string,
-    startLine: number,
-    entry: string
-): string {
-    if (fileMatchRegex.test(entry)) {
-        return modifyLineNumbers(entry, file, startLine);
-    } else if (IPythonMatchRegex.test(entry)) {
-        const ipythonReplaced = entry.replace(IPythonMatchRegex, fileDisplayName);
-        return modifyLineNumbers(ipythonReplaced, file, startLine);
-    }
-    return entry;
-}
-
-export function modifyTraceback(
-    file: string,
-    fileDisplayName: string,
-    startLine: number,
-    traceback: string[]
-): string[] {
-    if (file && file !== Identifiers.EmptyFileName) {
-        const escaped = _escapeRegExp(fileDisplayName);
-        const fileMatchRegex = new RegExp(`\\[.*?;32m${escaped}`);
-        return traceback.map(modifyTracebackEntry.bind(undefined, fileMatchRegex, file, fileDisplayName, startLine));
-    }
-    return traceback;
-}

src/client/datascience/types.ts

@@ -249,6 +249,7 @@ export interface INotebookExecutionLogger {
     preExecute(cell: ICell, silent: boolean): Promise<void>;
     postExecute(cell: ICell, silent: boolean): Promise<void>;
     onKernelRestarted(): void;
+    preHandleIOPub?(msg: KernelMessage.IIOPubMessage): KernelMessage.IIOPubMessage;
 }
 
 export const IGatherProvider = Symbol('IGatherProvider');