stdlib included in auto-imports

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

@@ -337,12 +337,19 @@ export class ImportResolver {
         return this._getThirdPartyTypeshedPath(this._configOptions.typeshedPath, unused);
     }
 
+    getStdlibModules = (execEnv: ExecutionEnvironment) => {
+        if (!this._stdlibModules) {
+            this._stdlibModules = this._buildStdlibCache(this.getTypeshedStdLibPath(execEnv), execEnv);
+        }
+        return this._stdlibModules;
+    };
+
     isStdlibModule(module: ImportedModuleDescriptor, execEnv: ExecutionEnvironment): boolean {
         if (!this._stdlibModules) {
             this._stdlibModules = this._buildStdlibCache(this.getTypeshedStdLibPath(execEnv), execEnv);
         }
 
-        return this._stdlibModules.has(module.nameParts.join('.'));
+        return this.getStdlibModules(execEnv).has(module.nameParts.join('.'));
     }
 
     getImportRoots(execEnv: ExecutionEnvironment, forLogging = false) {

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

@@ -982,6 +982,19 @@ export class Program {
         return program;
     }
 
+    loadStdlibModules = (sourceFileUri: Uri) => {
+        this._importResolver
+            ?.getStdlibModules(this._configOptions.findExecEnvironment(sourceFileUri))
+            .forEach((module) => {
+                if (!module.includes('.')) {
+                    this._lookUpImport({
+                        importingFileUri: sourceFileUri,
+                        nameParts: [module],
+                    });
+                }
+            });
+    };
+
     // Returns a value from 0 to 1 (or more) indicating how "full" the cache is
     // relative to some predetermined high-water mark. We'll compute this value
     // based on two easy-to-compute metrics: the number of entries in the type

packages/pyright-internal/src/common/extensibility.ts

@@ -90,6 +90,9 @@ export interface ProgramView {
     // See whether we can get rid of these methods
     handleMemoryHighUsage(): void;
     clone(): prog.Program;
+
+    // TODO: does this "mutate the program"? if so it should be moved somewhere else
+    loadStdlibModules: (sourceFileUri: Uri) => void;
 }
 
 // This exposes some APIs to mutate program. Unlike ProgramMutator, this will only mutate this program

packages/pyright-internal/src/languageService/completionProvider.ts

@@ -309,6 +309,7 @@ export class CompletionProvider {
             return null;
         }
 
+        this.program.loadStdlibModules(this.fileUri);
         const completionMap = this._getCompletions();
         return CompletionList.create(completionMap?.toArray());
     }

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

@@ -369,12 +369,12 @@ test('include literals in expression completion', async () => {
 //// from typing import TypedDict
 //// 
 //// class TestType(TypedDict):
-////     A: str
-////     B: int
+////     key_a: str
+////     key_b: int
 //// 
 //// var: TestType = {}
 //// 
-//// var[[|A/*marker*/|]]
+//// var[[|key_a/*marker*/|]]
     `;
 
     const state = parseAndGetTestState(code).state;
@@ -386,8 +386,8 @@ test('include literals in expression completion', async () => {
             completions: [
                 {
                     kind: CompletionItemKind.Constant,
-                    label: "'A'",
-                    textEdit: { range: state.getPositionRange('marker'), newText: "'A'" },
+                    label: "'key_a'",
+                    textEdit: { range: state.getPositionRange('marker'), newText: "'key_a'" },
                 },
             ],
         },
@@ -400,10 +400,10 @@ test('include literals in set key', async () => {
 //// from typing import TypedDict
 //// 
 //// class TestType(TypedDict):
-////     A: str
-////     B: int
+////     key_a: str
+////     key_b: int
 //// 
-//// var: TestType = { [|A/*marker*/|] }
+//// var: TestType = { [|key_a/*marker*/|] }
     `;
 
     const state = parseAndGetTestState(code).state;
@@ -415,8 +415,8 @@ test('include literals in set key', async () => {
             completions: [
                 {
                     kind: CompletionItemKind.Constant,
-                    label: "'A'",
-                    textEdit: { range: state.getPositionRange('marker'), newText: "'A'" },
+                    label: "'key_a'",
+                    textEdit: { range: state.getPositionRange('marker'), newText: "'key_a'" },
                 },
             ],
         },
@@ -429,10 +429,10 @@ test('include literals in dict key', async () => {
 //// from typing import TypedDict
 //// 
 //// class TestType(TypedDict):
-////     A: str
-////     B: int
+////     key_a: str
+////     key_b: int
 //// 
-//// var: TestType = { [|A/*marker*/|] : "hello" }
+//// var: TestType = { [|key_a/*marker*/|] : "hello" }
     `;
 
     const state = parseAndGetTestState(code).state;
@@ -444,8 +444,8 @@ test('include literals in dict key', async () => {
             completions: [
                 {
                     kind: CompletionItemKind.Constant,
-                    label: '"A"',
-                    textEdit: { range: state.getPositionRange('marker'), newText: '"A"' },
+                    label: '"key_a"',
+                    textEdit: { range: state.getPositionRange('marker'), newText: '"key_a"' },
                 },
             ],
         },

packages/pyright-internal/src/tests/fourslash/completions.builtinOverride.fourslash.ts

@@ -1,7 +1,7 @@
 /// <reference path="typings/fourslash.d.ts" />
 
 // @filename: test.py
-//// Cust[|/*marker1*/|]
+//// Custo[|/*marker1*/|]
 //// my_v[|/*marker2*/|]
 
 // @filename: __builtins__.pyi

packages/pyright-internal/src/tests/fourslash/completions.declNames.method.fourslash.ts

@@ -20,7 +20,7 @@
 //// def method(p:[|/*marker8*/|]):
 ////     pass
 ////
-//// def method(p, p2[|/*marker9*/|]):
+//// def method(p, p21234[|/*marker9*/|]):
 ////     pass
 //// def method(p, p2:[|/*marker10*/|]):
 ////     pass
@@ -37,7 +37,7 @@
 ////     def method(p:[|/*marker14*/|]):
 ////         pass
 ////
-////     def method(p, p2[|/*marker15*/|]):
+////     def method(p, p21234[|/*marker15*/|]):
 ////         pass
 ////     def method(p, p2:[|/*marker16*/|]):
 ////         pass