pull pylance (#2434)

Co-authored-by: Jake Bailey <5341706+jakebailey@users.noreply.github.com>
Co-authored-by: Heejae Chang <heejaechang@outlook.com>

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

@@ -10,7 +10,7 @@
 
 import type { Dirent } from 'fs';
 
-import { getOrAdd } from '../common/collectionUtils';
+import { flatten, getMapValues, getOrAdd } from '../common/collectionUtils';
 import { ConfigOptions, ExecutionEnvironment } from '../common/configOptions';
 import { FileSystem } from '../common/fileSystem';
 import { Host } from '../common/host';
@@ -1533,7 +1533,18 @@ export class ImportResolver {
                 typeshedPaths = [path];
             }
         } else {
-            typeshedPaths = this._getThirdPartyTypeshedPackagePaths(moduleDescriptor, execEnv, importFailureInfo);
+            typeshedPaths = this._getThirdPartyTypeshedPackagePaths(
+                moduleDescriptor,
+                execEnv,
+                importFailureInfo,
+                /*includeMatchOnly*/ false
+            );
+
+            const typeshedPathEx = this.getTypeshedPathEx(execEnv, importFailureInfo);
+            if (typeshedPathEx) {
+                typeshedPaths = typeshedPaths ?? [];
+                typeshedPaths.push(typeshedPathEx);
+            }
         }
 
         if (!typeshedPaths) {
@@ -1668,7 +1679,8 @@ export class ImportResolver {
     private _getThirdPartyTypeshedPackagePaths(
         moduleDescriptor: ImportedModuleDescriptor,
         execEnv: ExecutionEnvironment,
-        importFailureInfo: string[]
+        importFailureInfo: string[],
+        includeMatchOnly = true
     ): string[] | undefined {
         const typeshedPath = this._getThirdPartyTypeshedPath(execEnv, importFailureInfo);
 
@@ -1677,7 +1689,17 @@ export class ImportResolver {
         }
 
         const firstNamePart = moduleDescriptor.nameParts.length > 0 ? moduleDescriptor.nameParts[0] : '';
-        return this._cachedTypeshedThirdPartyPackagePaths!.get(firstNamePart);
+        if (includeMatchOnly) {
+            return this._cachedTypeshedThirdPartyPackagePaths!.get(firstNamePart);
+        }
+
+        if (firstNamePart) {
+            return flatten(
+                getMapValues(this._cachedTypeshedThirdPartyPackagePaths!, (k) => k.startsWith(firstNamePart))
+            );
+        }
+
+        return [];
     }
 
     private _getThirdPartyTypeshedPackageRoots(execEnv: ExecutionEnvironment, importFailureInfo: string[]) {

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

@@ -350,3 +350,14 @@ export function addIfUnique<T>(arr: T[], t: T): T[] {
     arr.push(t);
     return arr;
 }
+
+export function getMapValues<K, V>(m: Map<K, V>, predicate: (k: K, v: V) => boolean): V[] {
+    const values: V[] = [];
+    m.forEach((v, k) => {
+        if (predicate(k, v)) {
+            values.push(v);
+        }
+    });
+
+    return values;
+}

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

@@ -0,0 +1,9 @@
+/// <reference path="fourslash.ts" />
+
+// @filename: test.py
+//// from r/*marker*/
+
+// @ts-ignore
+await helper.verifyCompletion('included', 'markdown', {
+    marker: { completions: [{ label: 'requests', kind: Consts.CompletionItemKind.Module }] },
+});