Set isThirdPartyImport and isThirdPartyPyTypedPresent for interim…

… files (#6155)
microsoft · Oct 23, 2023 · 5d3d703 · 5d3d703
1 parent 6e2dcee
commit 5d3d703
Show file tree

Hide file tree

Showing 19 changed files with 217 additions and 34 deletions.
diff --git a/packages/pyright-internal/src/analyzer/importResolver.ts b/packages/pyright-internal/src/analyzer/importResolver.ts
@@ -38,6 +38,8 @@ import {
     tryStat,
 } from '../common/pathUtils';
 import { PythonVersion, versionFromString } from '../common/pythonVersion';
+import { ServiceProvider } from '../common/serviceProvider';
+import { ServiceKeys } from '../common/serviceProviderExtensions';
 import * as StringUtils from '../common/stringUtils';
 import { equateStringsCaseInsensitive } from '../common/stringUtils';
 import { isIdentifierChar, isIdentifierStartChar } from '../parser/characters';
@@ -48,8 +50,6 @@ import { PyTypedInfo, getPyTypedInfo } from './pyTypedUtils';
 import * as PythonPathUtils from './pythonPathUtils';
 import * as SymbolNameUtils from './symbolNameUtils';
 import { isDunderName } from './symbolNameUtils';
-import { ServiceProvider } from '../common/serviceProvider';
-import { ServiceKeys } from '../common/serviceProviderExtensions';
 
 export interface ImportedModuleDescriptor {
     leadingDots: number;
@@ -64,6 +64,10 @@ export interface ModuleNameAndType {
     isLocalTypingsFile: boolean;
 }
 
+export interface ModuleImportInfo extends ModuleNameAndType {
+    isThirdPartyPyTypedPresent: boolean;
+}
+
 export interface ModuleNameInfoFromPath {
     moduleName: string;
     containsInvalidCharacters?: boolean;
@@ -111,7 +115,7 @@ const allowPartialResolutionForThirdPartyPackages = false;
 export class ImportResolver {
     private _cachedPythonSearchPaths: { paths: string[]; failureInfo: string[] } | undefined;
     private _cachedImportResults = new Map<string | undefined, CachedImportResults>();
-    private _cachedModuleNameResults = new Map<string, Map<string, ModuleNameAndType>>();
+    private _cachedModuleNameResults = new Map<string, Map<string, ModuleImportInfo>>();
     private _cachedTypeshedRoot: string | undefined;
     private _cachedTypeshedStdLibPath: string | undefined;
     private _cachedTypeshedStdLibModuleVersions: Map<string, SupportedVersionRange> | undefined;
@@ -136,7 +140,7 @@ export class ImportResolver {
 
     invalidateCache() {
         this._cachedImportResults = new Map<string | undefined, CachedImportResults>();
-        this._cachedModuleNameResults = new Map<string, Map<string, ModuleNameAndType>>();
+        this._cachedModuleNameResults = new Map<string, Map<string, ModuleImportInfo>>();
         this.cachedParentImportResults.reset();
         this._stdlibModules = undefined;
 
@@ -311,7 +315,7 @@ export class ImportResolver {
     // In a sense, it's performing the inverse of resolveImport.
     getModuleNameForImport(filePath: string, execEnv: ExecutionEnvironment, allowInvalidModuleName = false) {
         // Cache results of the reverse of resolveImport as we cache resolveImport.
-        const cache = getOrAdd(this._cachedModuleNameResults, execEnv.root, () => new Map<string, ModuleNameAndType>());
+        const cache = getOrAdd(this._cachedModuleNameResults, execEnv.root, () => new Map<string, ModuleImportInfo>());
         return getOrAdd(cache, filePath, () => this._getModuleNameForImport(filePath, execEnv, allowInvalidModuleName));
     }
 
@@ -1046,10 +1050,11 @@ export class ImportResolver {
         filePath: string,
         execEnv: ExecutionEnvironment,
         allowInvalidModuleName: boolean
-    ): ModuleNameAndType {
+    ): ModuleImportInfo {
         let moduleName: string | undefined;
         let importType = ImportType.BuiltIn;
         let isLocalTypingsFile = false;
+        let isThirdPartyPyTypedPresent = false;
 
         const importFailureInfo: string[] = [];
 
@@ -1083,7 +1088,12 @@ export class ImportResolver {
                         []
                     )
                 ) {
-                    return { moduleName, importType, isLocalTypingsFile };
+                    return {
+                        moduleName,
+                        importType,
+                        isLocalTypingsFile,
+                        isThirdPartyPyTypedPresent,
+                    };
                 }
             }
         }
@@ -1197,16 +1207,48 @@ export class ImportResolver {
             }
         }
 
+        if (importType === ImportType.ThirdParty) {
+            const root = this.getParentImportResolutionRoot(filePath, execEnv.root);
+
+            // Go up directories one by one looking for a py.typed file.
+            let current = ensureTrailingDirectorySeparator(getDirectoryPath(filePath));
+            while (this._shouldWalkUp(current, root, execEnv)) {
+                if (this.fileExistsCached(combinePaths(current, 'py.typed'))) {
+                    const pyTypedInfo = getPyTypedInfo(this.fileSystem, current);
+                    if (pyTypedInfo && !pyTypedInfo.isPartiallyTyped) {
+                        isThirdPartyPyTypedPresent = true;
+                    }
+                    break;
+                }
+
+                let success;
+                [success, current] = this._tryWalkUp(current);
+                if (!success) {
+                    break;
+                }
+            }
+        }
+
         if (moduleName) {
-            return { moduleName, importType, isLocalTypingsFile };
+            return { moduleName, importType, isLocalTypingsFile, isThirdPartyPyTypedPresent };
         }
 
         if (allowInvalidModuleName && moduleNameWithInvalidCharacters) {
-            return { moduleName: moduleNameWithInvalidCharacters, importType, isLocalTypingsFile };
+            return {
+                moduleName: moduleNameWithInvalidCharacters,
+                importType,
+                isLocalTypingsFile,
+                isThirdPartyPyTypedPresent,
+            };
         }
 
         // We didn't find any module name.
-        return { moduleName: '', importType: ImportType.Local, isLocalTypingsFile };
+        return {
+            moduleName: '',
+            importType: ImportType.Local,
+            isLocalTypingsFile,
+            isThirdPartyPyTypedPresent,
+        };
     }
 
     private _invalidateFileSystemCache() {

diff --git a/packages/pyright-internal/src/analyzer/program.ts b/packages/pyright-internal/src/analyzer/program.ts
@@ -46,6 +46,7 @@ import { ImportResult, ImportType } from './importResult';
 import { getDocString } from './parseTreeUtils';
 import { Scope } from './scope';
 import { IPythonMode, SourceFile, SourceFileEditMode } from './sourceFile';
+import { SourceFileInfo } from './sourceFileInfo';
 import { createChainedByList, isUserCode, verifyNoCyclesInChainedFiles } from './sourceFileInfoUtils';
 import { SourceMapper } from './sourceMapper';
 import { Symbol } from './symbol';
@@ -55,7 +56,6 @@ import { createTypeEvaluatorWithTracker } from './typeEvaluatorWithTracker';
 import { PrintTypeFlags } from './typePrinter';
 import { TypeStubWriter } from './typeStubWriter';
 import { Type } from './types';
-import { SourceFileInfo } from './sourceFileInfo';
 
 const _maxImportDepth = 256;
 
@@ -1575,22 +1575,26 @@ export class Program {
     }
 
     private _createInterimFileInfo(filePath: string) {
-        const importName = this._getImportNameForFile(filePath);
+        const moduleImportInfo = this._importResolver.getModuleNameForImport(
+            filePath,
+            this._configOptions.getDefaultExecEnvironment(),
+            /* allowIllegalModuleName */ true
+        );
         const sourceFile = this._sourceFileFactory.createSourceFile(
             this.serviceProvider,
             filePath,
-            importName,
-            /* isThirdPartyImport */ false,
-            /* isInPyTypedPackage */ false,
+            moduleImportInfo.moduleName,
+            moduleImportInfo.importType === ImportType.ThirdParty,
+            moduleImportInfo.isThirdPartyPyTypedPresent,
             this._editModeTracker,
             this._console,
             this._logTracker
         );
         const sourceFileInfo = new SourceFileInfo(
             sourceFile,
             /* isTypeshedFile */ false,
-            /* isThirdPartyImport */ false,
-            /* isThirdPartyPyTypedPresent */ false,
+            moduleImportInfo.importType === ImportType.ThirdParty,
+            moduleImportInfo.isThirdPartyPyTypedPresent,
             this._editModeTracker
         );
 

diff --git a/packages/pyright-internal/src/tests/fourslash/completions.importInterimFile.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/completions.importInterimFile.fourslash.ts
@@ -0,0 +1,52 @@
+/// <reference path="fourslash.ts" />
+
+// @filename: test1.py
+//// import al[|/*marker1*/|]
+
+// @filename: altair/__init__.py
+// @library: true
+//// """module docstring"""
+////
+//// __all__ = [ "selection_interval" ]
+////
+//// from .vegalite import (
+////     selection,
+////     selection_interval
+//// )
+
+// @filename: altair/vegalite/__init__.py
+// @library: true
+//// def selection(): pass
+//// def selection_interval(): pass
+
+// @filename: altair/py.typed
+// @library: true
+//// # has to contain something for file to be written
+
+{
+    // Force interim file to be created
+    // @ts-ignore
+    await helper.verifyCompletion('included', 'markdown', {
+        marker1: {
+            completions: [
+                { label: 'altair', kind: Consts.CompletionItemKind.Module, documentation: 'module docstring' },
+            ],
+        },
+    });
+
+    helper.replace(helper.BOF, helper.getPosition('marker1'), 'import altair as alt\n\nalt.');
+
+    // @ts-ignore
+    await helper.verifyCompletion('included', 'markdown', {
+        marker1: {
+            completions: [{ label: 'selection_interval', kind: Consts.CompletionItemKind.Function }],
+        },
+    });
+
+    // @ts-ignore
+    await helper.verifyCompletion('excluded', 'markdown', {
+        marker1: {
+            completions: [{ label: 'selection', kind: Consts.CompletionItemKind.Function }],
+        },
+    });
+}
diff --git a/packages/pyright-internal/src/tests/fourslash/completions.wildcardimports.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/completions.wildcardimports.fourslash.ts
@@ -1,5 +1,10 @@
 /// <reference path="fourslash.ts" />
 
+// @filename: pyrightconfig.json
+//// {
+////   "useLibraryCodeForTypes": true
+//// }
+
 // @filename: lib1/definition.py
 // @library: true
 //// def func():

diff --git a/packages/pyright-internal/src/tests/fourslash/findDefinitions.classes.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/findDefinitions.classes.fourslash.ts
@@ -1,5 +1,10 @@
 /// <reference path="fourslash.ts" />
 
+// @filename: pyrightconfig.json
+//// {
+////   "useLibraryCodeForTypes": true
+//// }
+
 // @filename: testLib1/__init__.pyi
 // @library: true
 //// class C: ...

diff --git a/packages/pyright-internal/src/tests/fourslash/findDefinitions.fields.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/findDefinitions.fields.fourslash.ts
@@ -1,5 +1,10 @@
 /// <reference path="fourslash.ts" />
 
+// @filename: pyrightconfig.json
+//// {
+////   "useLibraryCodeForTypes": true
+//// }
+
 // @filename: testLib1/__init__.pyi
 // @library: true
 //// from typing import ClassVar

diff --git a/packages/pyright-internal/src/tests/fourslash/findDefinitions.functions.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/findDefinitions.functions.fourslash.ts
@@ -1,5 +1,10 @@
 /// <reference path="fourslash.ts" />
 
+// @filename: pyrightconfig.json
+//// {
+////   "useLibraryCodeForTypes": true
+//// }
+
 // @filename: testLib1/__init__.pyi
 // @library: true
 //// def C(): ...

diff --git a/packages/pyright-internal/src/tests/fourslash/findDefinitions.methods.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/findDefinitions.methods.fourslash.ts
@@ -1,5 +1,10 @@
 /// <reference path="fourslash.ts" />
 
+// @filename: pyrightconfig.json
+//// {
+////   "useLibraryCodeForTypes": true
+//// }
+
 // @filename: testLib1/__init__.pyi
 // @library: true
 //// from typing import overload

diff --git a/packages/pyright-internal/src/tests/fourslash/findDefinitions.parameters.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/findDefinitions.parameters.fourslash.ts
@@ -1,5 +1,10 @@
 /// <reference path="fourslash.ts" />
 
+// @filename: pyrightconfig.json
+//// {
+////   "useLibraryCodeForTypes": true
+//// }
+
 // @filename: testLib1/__init__.pyi
 // @library: true
 //// from typing import overload

diff --git a/packages/pyright-internal/src/tests/fourslash/findDefinitions.variables.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/findDefinitions.variables.fourslash.ts
@@ -1,5 +1,10 @@
 /// <reference path="fourslash.ts" />
 
+// @filename: pyrightconfig.json
+//// {
+////   "useLibraryCodeForTypes": true
+//// }
+
 // @filename: testLib1/__init__.pyi
 // @library: true
 //// C = ...

diff --git a/packages/pyright-internal/src/tests/fourslash/findDefinitions.wildcardimports.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/findDefinitions.wildcardimports.fourslash.ts
@@ -1,5 +1,10 @@
 /// <reference path="fourslash.ts" />
 
+// @filename: pyrightconfig.json
+//// {
+////   "useLibraryCodeForTypes": true
+//// }
+
 // @filename: lib1/definition.py
 // @library: true
 //// def [|func|]():

diff --git a/packages/pyright-internal/src/tests/fourslash/findTypeDefinitions.classes.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/findTypeDefinitions.classes.fourslash.ts
@@ -1,5 +1,10 @@
 /// <reference path="fourslash.ts" />
 
+// @filename: pyrightconfig.json
+//// {
+////   "useLibraryCodeForTypes": true
+//// }
+
 // @filename: testLib1/__init__.pyi
 // @library: true
 //// class [|C|]: ...

diff --git a/packages/pyright-internal/src/tests/fourslash/fourslash.ts b/packages/pyright-internal/src/tests/fourslash/fourslash.ts
@@ -233,7 +233,9 @@ declare namespace _ {
             predicate: (m: Marker | undefined, d: T | undefined, text: string) => boolean
         ): Range[];
         getPositionRange(markerString: string): PositionRange;
-        getPosition(markerString: string): Position;
+        getPosition(markerString: string): number;
+        get BOF(): number;
+        get EOF(): number;
         expandPositionRange(range: PositionRange, start: number, end: number): PositionRange;
         convertPositionRange(range: Range): PositionRange;
         convertPathToUri(path: string): string;
@@ -347,11 +349,11 @@ declare namespace _ {
             isUntitled?: boolean
         ): void;
 
+        replace(start: number, length: number, text: string): void;
+
         /* not tested yet
         paste(text: string): void;
-
         type(text: string): void;
-        replace(start: number, length: number, text: string): void;
         deleteChar(count: number): void;
         deleteLineRange(startIndex: number, endIndexInclusive: number): void;
         deleteCharBehindMarker(count: number): void;

diff --git a/packages/pyright-internal/src/tests/fourslash/hover.class.docString.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/hover.class.docString.fourslash.ts
@@ -1,5 +1,10 @@
 /// <reference path="fourslash.ts" />
 
+// @filename: pyrightconfig.json
+//// {
+////   "useLibraryCodeForTypes": true
+//// }
+
 // @filename: test.py
 //// import lib
 ////

diff --git a/packages/pyright-internal/src/tests/fourslash/hover.docFromSrc.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/hover.docFromSrc.fourslash.ts
@@ -1,5 +1,10 @@
 /// <reference path="fourslash.ts" />
 
+// @filename: pyrightconfig.json
+//// {
+////   "useLibraryCodeForTypes": true
+//// }
+
 // @filename: module1.py
 //// '''module1 docs'''
 ////

diff --git a/...over.docFromScr.stringFormat.fourslash.ts → ...over.docFromSrc.stringFormat.fourslash.ts b/...over.docFromScr.stringFormat.fourslash.ts → ...over.docFromSrc.stringFormat.fourslash.ts
diff --git a/packages/pyright-internal/src/tests/fourslash/hover.wildcardimports.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/hover.wildcardimports.fourslash.ts
@@ -1,5 +1,10 @@
 /// <reference path="fourslash.ts" />
 
+// @filename: pyrightconfig.json
+//// {
+////   "useLibraryCodeForTypes": true
+//// }
+
 // @filename: lib1/definition.py
 // @library: true
 //// def func():