unit tests for jupyter cli support

packages/pyright-internal/src/analyzer/sourceFile.ts

@@ -459,7 +459,7 @@ export class SourceFile {
         // that of the previous contents.
         try {
             // Read the file's contents.
-            const uri = this._getRealUri();
+            const uri = this.getRealUri();
             if (this.fileSystem.existsSync(uri)) {
                 const fileContents = this.fileSystem.readFileSync(uri, 'utf8');
 
@@ -554,7 +554,7 @@ export class SourceFile {
             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)]
+        const source = getIPythonCells(this.fileSystem, this.getRealUri(), this._console)?.[getCellIndex(this._uri)]
             .source;
         return typeof source === 'string' ? source : source?.join('');
     }
@@ -1004,6 +1004,15 @@ export class SourceFile {
         });
     }
 
+    /**
+     * 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.
+     */
+    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('');
+
     test_enableIPythonMode(enable: boolean) {
         this._ipythonMode = enable ? IPythonMode.CellDocs : IPythonMode.None;
     }
@@ -1016,15 +1025,6 @@ 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

packages/pyright-internal/src/tests/notebooks.test.ts

@@ -0,0 +1,26 @@
+import { tExpect } from 'typed-jest-expect';
+import { DiagnosticRule } from '../common/diagnosticRules';
+import { typeAnalyzeSampleFiles, validateResultsButBased } from './testUtils';
+
+test('symbol from previous cell', () => {
+    const analysisResults = typeAnalyzeSampleFiles(['notebook.ipynb']);
+    tExpect(analysisResults.length).toStrictEqual(2);
+    validateResultsButBased(analysisResults[0], {
+        errors: [],
+    });
+    validateResultsButBased(analysisResults[1], {
+        errors: [
+            {
+                code: DiagnosticRule.reportAssignmentType,
+                line: 0,
+                //TODO: get rid of these stupid fake spaces
+                message: 'Type "int" is not assignable to declared type "str"\n  "int" is not assignable to "str"',
+            },
+        ],
+    });
+});
+
+test('non-python cell', () => {
+    const analysisResults = typeAnalyzeSampleFiles(['notebook2.ipynb']);
+    tExpect(analysisResults.length).toStrictEqual(0);
+});

packages/pyright-internal/src/tests/samples/notebook.ipynb

@@ -0,0 +1,29 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "foo: int = 1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "bar: str = foo"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

packages/pyright-internal/src/tests/samples/notebook2.ipynb

@@ -0,0 +1,24 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "vscode": {
+     "languageId": "plaintext"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "asdf"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

packages/pyright-internal/src/tests/testUtils.ts

@@ -33,6 +33,7 @@ import { TypeInlayHintsItemType, TypeInlayHintsWalker } from '../analyzer/typeIn
 import { Range } from 'vscode-languageserver-types';
 import { ServiceProvider } from '../common/serviceProvider';
 import { InlayHintSettings } from '../workspaceFactory';
+import { getCellIndex } from '../analyzer/sourceFile';
 
 // This is a bit gross, but it's necessary to allow the fallback typeshed
 // directory to be located when running within the jest environment. This
@@ -42,6 +43,7 @@ import { InlayHintSettings } from '../workspaceFactory';
 
 export interface FileAnalysisResult {
     fileUri: Uri;
+    cell: number | undefined;
     parseResults?: ParseFileResults | undefined;
     errors: Diagnostic[];
     warnings: Diagnostic[];
@@ -180,12 +182,21 @@ export function getAnalysisResults(
         // specifying a timeout, it should complete the first time.
     }
 
-    const sourceFiles = fileUris.map((filePath) => program.getSourceFile(filePath));
+    const sourceFiles = fileUris
+        .flatMap((filePath) =>
+            program
+                .getSourceFileInfoList()
+                // find notebook cells
+                .filter((sourceFileInfo) => sourceFileInfo.sourceFile.getRealUri().equals(filePath))
+        )
+        .map((sourceFile) => sourceFile.sourceFile);
     return sourceFiles.map((sourceFile, index) => {
         if (sourceFile) {
             const diagnostics = sourceFile.getDiagnostics(configOptions) || [];
+            const fileUri = sourceFile.getUri();
             const analysisResult: FileAnalysisResult = {
-                fileUri: sourceFile.getUri(),
+                fileUri,
+                cell: getCellIndex(fileUri),
                 parseResults: sourceFile.getParseResults(),
                 errors: diagnostics.filter((diag) => diag.category === DiagnosticCategory.Error),
                 warnings: diagnostics.filter((diag) => diag.category === DiagnosticCategory.Warning),
@@ -251,11 +262,20 @@ export type ExpectedResults = {
     [key in Exclude<keyof FileAnalysisResult, 'fileUri' | 'parseResults'>]?: ExpectedResult[];
 };
 
-export const validateResultsButBased = (allResults: FileAnalysisResult[], expectedResults: ExpectedResults) => {
-    assert.strictEqual(allResults.length, 1);
-    const result = allResults[0];
+export const validateResultsButBased = (
+    allResults: FileAnalysisResult | FileAnalysisResult[],
+    expectedResults: ExpectedResults
+) => {
+    let result: FileAnalysisResult;
+    if (Array.isArray(allResults)) {
+        assert.strictEqual(allResults.length, 1);
+        result = allResults[0];
+    } else {
+        result = allResults;
+    }
     for (const [diagnosticType] of entries(result)) {
-        if (diagnosticType === 'fileUri' || diagnosticType === 'parseResults') {
+        // TODO: this is gross, also update it so that you can specify expected notebook cells
+        if (diagnosticType === 'fileUri' || diagnosticType === 'parseResults' || diagnosticType === 'cell') {
             continue;
         }
         const actualResult = result[diagnosticType].map(