automatically import IPython.display.display in notebooks to match …

…runtime and pylance
DetachHead · Feb 23, 2025 · 77e8fa4 · 77e8fa4
1 parent aea3ee0
commit 77e8fa4
Show file tree

Hide file tree

Showing 6 changed files with 93 additions and 28 deletions.
diff --git a/packages/pyright-internal/src/analyzer/binder.ts b/packages/pyright-internal/src/analyzer/binder.ts
@@ -123,6 +123,7 @@ import * as ParseTreeUtils from './parseTreeUtils';
 import { ParseTreeWalker } from './parseTreeWalker';
 import { moduleIsInList } from './pythonPathUtils';
 import { NameBindingType, Scope, ScopeType } from './scope';
+import { IPythonMode } from './sourceFile';
 import * as StaticExpressions from './staticExpressions';
 import { Symbol, SymbolFlags, indeterminateSymbolId } from './symbol';
 import { isConstantName, isPrivateName, isPrivateOrProtectedName } from './symbolNameUtils';
@@ -294,6 +295,11 @@ export class Binder extends ParseTreeWalker {
                 this._addImplicitSymbolToCurrentScope('__annotations__', node, 'Dict[str, Any]');
                 this._addImplicitSymbolToCurrentScope('__builtins__', node, 'Any');
                 this._addImplicitSymbolToCurrentScope('__doc__', node, 'str | None');
+                if (this._fileInfo.ipythonMode === IPythonMode.CellDocs) {
+                    // this function is automatically available globally inside notebooks.
+                    // https://ipython.readthedocs.io/en/stable/api/generated/IPython.display.html#IPython.display.display
+                    this._addImplicitSymbolToCurrentScope('display', node, 'IPython.display.display');
+                }
 
                 // Create a start node for the module.
                 this._currentFlowNode = this._createStartFlowNode();

diff --git a/packages/pyright-internal/src/analyzer/declaration.ts b/packages/pyright-internal/src/analyzer/declaration.ts
@@ -46,7 +46,15 @@ export const enum DeclarationType {
     Alias,
 }
 
-export type IntrinsicType = 'Any' | 'str' | 'str | None' | 'int' | 'Iterable[str]' | 'type[self]' | 'Dict[str, Any]';
+export type IntrinsicType =
+    | 'Any'
+    | 'str'
+    | 'str | None'
+    | 'int'
+    | 'Iterable[str]'
+    | 'type[self]'
+    | 'Dict[str, Any]'
+    | 'IPython.display.display';
 
 export interface DeclarationBase {
     // Category of this symbol (function, variable, etc.).

diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts
@@ -22299,39 +22299,52 @@ export function createTypeEvaluator(
                             : UnknownType.create(),
                     };
                 }
-
-                const strType = getBuiltInObject(declaration.node, 'str');
-                const intType = getBuiltInObject(declaration.node, 'int');
-                if (isClassInstance(intType) && isClassInstance(strType)) {
-                    if (declaration.intrinsicType === 'str') {
-                        return { type: strType };
+                if (declaration.intrinsicType === 'IPython.display.display') {
+                    const fileInfo = AnalyzerNodeInfo.getFileInfo(declaration.node);
+                    const lookupResult = fileInfo.importLookup({
+                        importingFileUri: fileInfo.fileUri,
+                        nameParts: ['IPython', 'display'],
+                    });
+                    if (lookupResult) {
+                        const symbol = lookupResult.symbolTable.get('display');
+                        if (symbol) {
+                            return { type: getEffectiveTypeOfSymbol(symbol) };
+                        }
                     }
+                } else {
+                    const strType = getBuiltInObject(declaration.node, 'str');
+                    const intType = getBuiltInObject(declaration.node, 'int');
+                    if (isClassInstance(intType) && isClassInstance(strType)) {
+                        if (declaration.intrinsicType === 'str') {
+                            return { type: strType };
+                        }
 
-                    if (declaration.intrinsicType === 'str | None') {
-                        return { type: combineTypes([strType, getNoneType()]) };
-                    }
+                        if (declaration.intrinsicType === 'str | None') {
+                            return { type: combineTypes([strType, getNoneType()]) };
+                        }
 
-                    if (declaration.intrinsicType === 'int') {
-                        return { type: intType };
-                    }
+                        if (declaration.intrinsicType === 'int') {
+                            return { type: intType };
+                        }
 
-                    if (declaration.intrinsicType === 'Iterable[str]') {
-                        const iterableType = getBuiltInType(declaration.node, 'Iterable');
-                        if (isInstantiableClass(iterableType)) {
-                            return {
-                                type: ClassType.cloneAsInstance(ClassType.specialize(iterableType, [strType])),
-                            };
+                        if (declaration.intrinsicType === 'Iterable[str]') {
+                            const iterableType = getBuiltInType(declaration.node, 'Iterable');
+                            if (isInstantiableClass(iterableType)) {
+                                return {
+                                    type: ClassType.cloneAsInstance(ClassType.specialize(iterableType, [strType])),
+                                };
+                            }
                         }
-                    }
 
-                    if (declaration.intrinsicType === 'Dict[str, Any]') {
-                        const dictType = getBuiltInType(declaration.node, 'dict');
-                        if (isInstantiableClass(dictType)) {
-                            return {
-                                type: ClassType.cloneAsInstance(
-                                    ClassType.specialize(dictType, [strType, AnyType.create()])
-                                ),
-                            };
+                        if (declaration.intrinsicType === 'Dict[str, Any]') {
+                            const dictType = getBuiltInType(declaration.node, 'dict');
+                            if (isInstantiableClass(dictType)) {
+                                return {
+                                    type: ClassType.cloneAsInstance(
+                                        ClassType.specialize(dictType, [strType, AnyType.create()])
+                                    ),
+                                };
+                            }
                         }
                     }
                 }

diff --git a/packages/pyright-internal/src/tests/notebooks.test.ts b/packages/pyright-internal/src/tests/notebooks.test.ts
@@ -34,3 +34,10 @@ test('invalid notebook file', () => {
         /failed to parse jupyter notebook .* - SyntaxError: Unexpected token .*/
     );
 });
+
+test('IPython.display.display automatically imported', () => {
+    const analysisResults = typeAnalyzeSampleFiles(['ipython_display_import/notebook.ipynb']);
+    validateResultsButBased(analysisResults, {
+        errors: [{ code: DiagnosticRule.reportUndefinedVariable, line: 4 }],
+    });
+});
diff --git a/packages/pyright-internal/src/tests/samples/ipython_display_import/IPython/display.pyi b/packages/pyright-internal/src/tests/samples/ipython_display_import/IPython/display.pyi
@@ -0,0 +1,7 @@
+"""
+fake stub for `IPython.display` so we don't have to use the real one and the test doesn't have to run inside a venv.
+
+see https://github.com/DetachHead/basedpyright/issues/994
+"""
+def display() -> None: ...
+def display_html() -> None: ...
diff --git a/packages/pyright-internal/src/tests/samples/ipython_display_import/notebook.ipynb b/packages/pyright-internal/src/tests/samples/ipython_display_import/notebook.ipynb
@@ -0,0 +1,24 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Callable, assert_type\n",
+    "\n",
+    "assert_type(display, Callable[[], None])\n",
+    "# ensure that only display is imported, not the whole module\n",
+    "display_html"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}