diff --git a/package-lock.json b/package-lock.json index eded31422..83514ee78 100644 --- a/package-lock.json +++ b/package-lock.json @@ -2076,6 +2076,30 @@ "@jridgewell/sourcemap-codec": "^1.4.14" } }, + "node_modules/@jupyterlab/nbformat": { + "version": "4.3.5", + "resolved": "https://registry.npmjs.org/@jupyterlab/nbformat/-/nbformat-4.3.5.tgz", + "integrity": "sha512-rTge5VNBKoPbTzwMagbZ8KgQa3R805eqPDVVPQhXb0F9ND3vfEbINbZOgz+6/5b0ABlbRnOYWWvTR4iU8cD77w==", + "license": "BSD-3-Clause", + "dependencies": { + "@lumino/coreutils": "^2.2.0" + } + }, + "node_modules/@lumino/algorithm": { + "version": "2.0.2", + "resolved": "https://registry.npmjs.org/@lumino/algorithm/-/algorithm-2.0.2.tgz", + "integrity": "sha512-cI8yJ2+QK1yM5ZRU3Kuaw9fJ/64JEDZEwWWp7+U0cd/mvcZ44BGdJJ29w+tIet1QXxPAvnsUleWyQ5qm4qUouA==", + "license": "BSD-3-Clause" + }, + "node_modules/@lumino/coreutils": { + "version": "2.2.0", + "resolved": "https://registry.npmjs.org/@lumino/coreutils/-/coreutils-2.2.0.tgz", + "integrity": "sha512-x5wnQ/GjWBayJ6vXVaUi6+Q6ETDdcUiH9eSfpRZFbgMQyyM6pi6baKqJBK2CHkCc/YbAEl6ipApTgm3KOJ/I3g==", + "license": "BSD-3-Clause", + "dependencies": { + "@lumino/algorithm": "^2.0.2" + } + }, "node_modules/@nodelib/fs.scandir": { "version": "2.1.5", "resolved": "https://registry.npmjs.org/@nodelib/fs.scandir/-/fs.scandir-2.1.5.tgz", @@ -14717,6 +14741,7 @@ "license": "MIT", "dependencies": { "@actions/core": "^1.10.1", + "@jupyterlab/nbformat": "^4.3.5", "@yarnpkg/fslib": "2.10.4", "@yarnpkg/libzip": "2.3.0", "chalk": "^4.1.2", diff --git a/packages/pyright-internal/package.json b/packages/pyright-internal/package.json index 9e911936b..f21568663 100644 --- a/packages/pyright-internal/package.json +++ b/packages/pyright-internal/package.json @@ -20,6 +20,7 @@ }, "dependencies": { "@actions/core": "^1.10.1", + "@jupyterlab/nbformat": "^4.3.5", "@yarnpkg/fslib": "2.10.4", "@yarnpkg/libzip": "2.3.0", "chalk": "^4.1.2", diff --git a/packages/pyright-internal/src/analyzer/program.ts b/packages/pyright-internal/src/analyzer/program.ts index c001ea395..44c32f2d5 100644 --- a/packages/pyright-internal/src/analyzer/program.ts +++ b/packages/pyright-internal/src/analyzer/program.ts @@ -40,7 +40,7 @@ import { ImportResult, ImportType } from './importResult'; import { getDocString } from './parseTreeUtils'; import { ISourceFileFactory } from './programTypes'; import { Scope } from './scope'; -import { IPythonMode, SourceFile } from './sourceFile'; +import { getIPythonCells, IPythonMode, SourceFile } from './sourceFile'; import { SourceFileInfo } from './sourceFileInfo'; import { createChainedByList, isUserCode, verifyNoCyclesInChainedFiles } from './sourceFileInfoUtils'; import { SourceMapper } from './sourceMapper'; @@ -332,47 +332,73 @@ export class Program { return fileInfo; } - addTrackedFile( - fileUri: Uri, - isThirdPartyImport = false, - isInPyTypedPackage = false, - isTypeshedFile = false - ): SourceFile { - let sourceFileInfo = this.getSourceFileInfo(fileUri); - const moduleImportInfo = this._getModuleImportInfoForFile(fileUri); - const importName = moduleImportInfo.moduleName; - - if (sourceFileInfo) { - // The module name may have changed based on updates to the - // search paths, so update it here. - sourceFileInfo.sourceFile.setModuleName(importName); - sourceFileInfo.isTracked = true; - return sourceFileInfo.sourceFile; + addTrackedFile(fileUri: Uri, isThirdPartyImport = false, isInPyTypedPackage = false, isTypeshedFile = false) { + const cells = getIPythonCells(this.fileSystem, fileUri, this.console); + const sourceFileInfos: SourceFileInfo[] = []; + if (cells) { + cells.forEach((_, index) => { + const cellUri = fileUri.withFragment(index.toString()); + const importName = this._getImportNameForNewSourceFile(cellUri); + if (!importName) { + return; // continue + } + const sourceFile = this._sourceFileFactory.createSourceFile( + this.serviceProvider, + cellUri, + importName, + isThirdPartyImport, + isInPyTypedPackage, + this.baselineHandler, + this._editModeTracker, + this._console, + this._logTracker, + IPythonMode.CellDocs + ); + sourceFileInfos.push( + new SourceFileInfo( + sourceFile, + isTypeshedFile, + isThirdPartyImport, + isInPyTypedPackage, + this._editModeTracker, + { + isTracked: true, + chainedSourceFile: sourceFileInfos[index - 1], + } + ) + ); + }); + } else { + const importName = this._getImportNameForNewSourceFile(fileUri); + if (!importName) { + return; + } + const sourceFile = this._sourceFileFactory.createSourceFile( + this.serviceProvider, + fileUri, + importName, + isThirdPartyImport, + isInPyTypedPackage, + this.baselineHandler, + this._editModeTracker, + this._console, + this._logTracker + ); + sourceFileInfos.push( + new SourceFileInfo( + sourceFile, + isTypeshedFile, + isThirdPartyImport, + isInPyTypedPackage, + this._editModeTracker, + { + isTracked: true, + } + ) + ); } - const sourceFile = this._sourceFileFactory.createSourceFile( - this.serviceProvider, - fileUri, - importName, - isThirdPartyImport, - isInPyTypedPackage, - this.baselineHandler, - this._editModeTracker, - this._console, - this._logTracker - ); - sourceFileInfo = new SourceFileInfo( - sourceFile, - isTypeshedFile, - isThirdPartyImport, - isInPyTypedPackage, - this._editModeTracker, - { - isTracked: true, - } - ); - this._addToSourceFileListAndMap(sourceFileInfo); - return sourceFile; + sourceFileInfos.forEach((sourceFileInfo) => this._addToSourceFileListAndMap(sourceFileInfo)); } setFileOpened(fileUri: Uri, version: number | null, contents: string, options?: OpenFileOptions) { @@ -1027,6 +1053,24 @@ export class Program { this.serviceProvider.tryGet(ServiceKeys.stateMutationListeners)?.forEach((l) => l.onClearCache?.()); } + /** + * @returns `undefined` if the source file is already tracked + */ + private _getImportNameForNewSourceFile = (fileUri: Uri): string | undefined => { + const sourceFileInfo = this.getSourceFileInfo(fileUri); + const moduleImportInfo = this._getModuleImportInfoForFile(fileUri); + const importName = moduleImportInfo.moduleName; + + if (sourceFileInfo) { + // The module name may have changed based on updates to the + // search paths, so update it here. + sourceFileInfo.sourceFile.setModuleName(importName); + sourceFileInfo.isTracked = true; + return undefined; + } + return importName; + }; + private _handleMemoryHighUsage() { const cacheUsage = this._cacheManager.getCacheUsage(); const usedHeapRatio = this._cacheManager.getUsedHeapRatio( diff --git a/packages/pyright-internal/src/analyzer/sourceFile.ts b/packages/pyright-internal/src/analyzer/sourceFile.ts index 8708981bd..69cb7bcf6 100644 --- a/packages/pyright-internal/src/analyzer/sourceFile.ts +++ b/packages/pyright-internal/src/analyzer/sourceFile.ts @@ -54,6 +54,7 @@ import { SymbolTable } from './symbol'; import { TestWalker } from './testWalker'; import { TypeEvaluator } from './typeEvaluatorTypes'; import { BaselineHandler } from '../baseline'; +import { INotebookContent } from '@jupyterlab/nbformat'; // Limit the number of import cycles tracked per source file. const _maxImportCyclesPerFile = 4; @@ -77,6 +78,60 @@ export enum IPythonMode { CellDocs, } +/** + * gets the content of a file. if it's a notebook, the whole notebook json is returned as it appears in the file + */ +const getFileContent = (fileSystem: FileSystem, uri: Uri, console: ConsoleInterface): string | undefined => { + try { + // Check the file's length before attempting to read its full contents. + const fileStat = fileSystem.statSync(uri); + if (fileStat.size > maxSourceFileSize) { + console.error( + `File length of "${uri}" is ${fileStat.size} ` + + `which exceeds the maximum supported file size of ${maxSourceFileSize}` + ); + throw new Error('File larger than max'); + } + + return fileSystem.readFileSync(uri, 'utf8'); + } catch (error) { + return undefined; + } +}; + +/** + * if this is a notebook, gets the cells in this file that contain python code. + * note that each {@link SourceFile} is an individual cell, but this function returns + * all of them + * @returns `undefined` if not a notebook + */ +export const getIPythonCells = (fileSystem: FileSystem, uri: Uri, console: ConsoleInterface) => { + if (!uri.hasExtension('.ipynb')) { + return undefined; + } + const fileContent = getFileContent(fileSystem, uri, console); + if (!fileContent) { + return undefined; + } + const parsedNotebook = JSON.parse(fileContent) as INotebookContent; + return parsedNotebook.cells.filter( + (cell) => + cell.cell_type === 'code' && + // i guess there's no standard way to specify the language of individual cells? so we check the metadata vscode adds for + // cells that have a different language to the notebook's language. idk if this is supported in any other editor tho + (typeof cell.metadata.vscode !== 'object' || + cell.metadata.vscode === null || + !('languageId' in cell.metadata.vscode) || + cell.metadata.vscode.languageId) + ); +}; + +/** + * gets the cell index from a notebook uri. must have a fragment containing the index. not used by the language server because + * it has its own fake notebook uri format + */ +export const getCellIndex = (uri: Uri): number => Number(uri.fragment); + // A monotonically increasing number used to create unique file IDs. let nextUniqueFileId = 1; @@ -403,8 +458,9 @@ export class SourceFile { // that of the previous contents. try { // Read the file's contents. - if (this.fileSystem.existsSync(this._uri)) { - const fileContents = this.fileSystem.readFileSync(this._uri, 'utf8'); + const uri = this._getRealUri(); + if (this.fileSystem.existsSync(uri)) { + const fileContents = this.fileSystem.readFileSync(uri, 'utf8'); if (fileContents.length !== this._writableData.lastFileContentLength) { return true; @@ -483,29 +539,23 @@ export class SourceFile { return this._writableData.clientDocumentContents; } + /** + * gets the content of the source file. if it's a notebook, the content of this source file's {@link _ipythonCellIndex} is returned + */ getFileContent(): string | undefined { - // Get current buffer content if the file is opened. - const openFileContent = this.getOpenFileContents(); - if (openFileContent !== undefined) { - return openFileContent; - } - - // Otherwise, get content from file system. - try { - // Check the file's length before attempting to read its full contents. - const fileStat = this.fileSystem.statSync(this._uri); - if (fileStat.size > maxSourceFileSize) { - this._console.error( - `File length of "${this._uri}" is ${fileStat.size} ` + - `which exceeds the maximum supported file size of ${maxSourceFileSize}` - ); - throw new Error('File larger than max'); + if (this._ipythonMode === IPythonMode.None) { + // Get current buffer content if the file is opened. + const openFileContent = this.getOpenFileContents(); + if (openFileContent !== undefined) { + return openFileContent; } - - return this.fileSystem.readFileSync(this._uri, 'utf8'); - } catch (error) { - return undefined; + // Otherwise, get content from file system. + return getFileContent(this.fileSystem, this._uri, this._console); } + //TODO: this isnt ideal because it re-reads the file for each cell which is unnecessary + const source = getIPythonCells(this.fileSystem, this._getRealUri(), this._console)?.[getCellIndex(this._uri)] + .source; + return typeof source === 'string' ? source : source?.join(''); } setClientVersion(version: number | null, contents: string): void { @@ -673,7 +723,7 @@ export class SourceFile { configOptions, this._uri, fileContents!, - this._ipythonMode, + this.getIPythonMode(), diagSink ); @@ -965,6 +1015,15 @@ export class SourceFile { return new TextRangeDiagnosticSink(lines); } + /** + * if this source file represents an ipython cell and we're in the cli instead of the language server, + * we need to create a fake uri to distinguish between them. + */ + private _getRealUri = () => + // when using the language server, a fake uri is assigned automatically and this method *should* never get called + // because we don't read the files from the disk directly, but we check this anyway just in case + this._writableData.clientDocumentContents ? this._uri : this._uri.withFragment(''); + // Creates a short string that can be used to uniquely identify // this file from all other files. It is used in the type evaluator // to distinguish between types that are defined in different files diff --git a/packages/pyright-internal/src/pyright.ts b/packages/pyright-internal/src/pyright.ts index dc788a80e..391cab140 100644 --- a/packages/pyright-internal/src/pyright.ts +++ b/packages/pyright-internal/src/pyright.ts @@ -22,7 +22,7 @@ import { AnalysisResults } from './analyzer/analysis'; import { PackageTypeReport, TypeKnownStatus } from './analyzer/packageTypeReport'; import { PackageTypeVerifier } from './analyzer/packageTypeVerifier'; import { AnalyzerService } from './analyzer/service'; -import { maxSourceFileSize } from './analyzer/sourceFile'; +import { getCellIndex, maxSourceFileSize } from './analyzer/sourceFile'; import { SourceFileInfo } from './analyzer/sourceFileInfo'; import { initializeDependencies } from './common/asyncInitialization'; import { ChokidarFileWatcherProvider } from './common/chokidarFileWatcherProvider'; @@ -124,6 +124,7 @@ interface PyrightPublicSymbolReport { // The schema for this object is publicly documented. Do not change it. interface PyrightJsonDiagnostic { file: string; + cell: number | undefined; severity: SeverityLevel; message: string; range?: Range | undefined; @@ -960,7 +961,7 @@ function buildTypeCompletenessReport( // Add the general diagnostics. completenessReport.generalDiagnostics.forEach((diag) => { - const jsonDiag = convertDiagnosticToJson('', diag); + const jsonDiag = convertDiagnosticToJson(undefined, diag); if (isDiagnosticIncluded(jsonDiag.severity, minSeverityLevel)) { report.generalDiagnostics.push(jsonDiag); } @@ -1007,7 +1008,7 @@ function buildTypeCompletenessReport( // Convert and filter the diagnostics. symbol.diagnostics.forEach((diag) => { - const jsonDiag = convertDiagnosticToJson(diag.uri.getFilePath(), diag.diagnostic); + const jsonDiag = convertDiagnosticToJson(diag.uri, diag.diagnostic); if (isDiagnosticIncluded(jsonDiag.severity, minSeverityLevel)) { diagnostics.push(jsonDiag); } @@ -1232,7 +1233,7 @@ function reportDiagnosticsAsJsonWithoutLogging( diag.category === DiagnosticCategory.Warning || diag.category === DiagnosticCategory.Information ) { - const jsonDiag = convertDiagnosticToJson(fileDiag.fileUri.getFilePath(), diag); + const jsonDiag = convertDiagnosticToJson(fileDiag.fileUri, diag); if (isDiagnosticIncluded(jsonDiag.severity, minSeverityLevel)) { report.generalDiagnostics.push(jsonDiag); } @@ -1299,9 +1300,10 @@ function convertDiagnosticCategoryToSeverity(category: DiagnosticCategory): Seve } } -function convertDiagnosticToJson(filePath: string, diag: Diagnostic): PyrightJsonDiagnostic { +function convertDiagnosticToJson(fileUri: Uri | undefined, diag: Diagnostic): PyrightJsonDiagnostic { return { - file: filePath, + file: fileUri?.getFilePath() ?? '', + cell: fileUri?.fragment ? getCellIndex(fileUri) : undefined, severity: convertDiagnosticCategoryToSeverity(diag.category), message: diag.message, range: isEmptyRange(diag.range) ? undefined : diag.range, @@ -1337,9 +1339,15 @@ function reportDiagnosticsAsText( ); if (fileErrorsAndWarnings.length > 0) { - console.info(`${fileDiagnostics.fileUri.toUserVisibleString()}`); - fileErrorsAndWarnings.forEach((diag) => { - const jsonDiag = convertDiagnosticToJson(fileDiagnostics.fileUri.getFilePath(), diag); + fileErrorsAndWarnings.forEach((diag, index) => { + const jsonDiag = convertDiagnosticToJson(fileDiagnostics.fileUri, diag); + if (index === 0) { + // only log this once per file. this is only in the for loop because we need to get the cell index from one of the diagnostics + console.info( + fileDiagnostics.fileUri.toUserVisibleString() + + (jsonDiag.cell ? ` - cell ${jsonDiag.cell + 1}` : '') + ); + } logDiagnosticToConsole(jsonDiag); if (diag.category === DiagnosticCategory.Error) { @@ -1446,6 +1454,9 @@ function logDiagnosticToConsole(diag: PyrightJsonDiagnostic, prefix = ' ') { if (diag.file) { message += `${diag.file}:`; } + if (diag.cell !== undefined) { + message += chalk.yellow(`${diag.cell + 1}`) + ':'; + } if (diag.range && !isEmptyRange(diag.range)) { message += chalk.yellow(`${diag.range.start.line + 1}`) +