Sync from Pylance (#2350)

Co-authored-by: HeeJae Chang <hechang@microsoft.com>

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

@@ -57,7 +57,7 @@ export function resolveAliasDeclaration(
         }
 
         let lookupResult: ImportLookupResult | undefined;
-        if (curDeclaration.path) {
+        if (curDeclaration.path && curDeclaration.loadSymbolsFromPath) {
             lookupResult = importLookup(curDeclaration.path);
         }
 

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

@@ -1559,6 +1559,7 @@ export class Binder extends ParseTreeWalker {
                                     type: DeclarationType.Alias,
                                     node,
                                     path: resolvedPath,
+                                    loadSymbolsFromPath: true,
                                     range: getEmptyRange(),
                                     usesLocalName: false,
                                     symbolName: name,
@@ -1579,6 +1580,7 @@ export class Binder extends ParseTreeWalker {
                                             type: DeclarationType.Alias,
                                             node,
                                             path: implicitImport.path,
+                                            loadSymbolsFromPath: true,
                                             range: getEmptyRange(),
                                             usesLocalName: false,
                                             moduleName: this._fileInfo.moduleName,
@@ -1588,6 +1590,7 @@ export class Binder extends ParseTreeWalker {
                                             type: DeclarationType.Alias,
                                             node,
                                             path: resolvedPath,
+                                            loadSymbolsFromPath: true,
                                             usesLocalName: false,
                                             symbolName: name,
                                             submoduleFallback,
@@ -1651,11 +1654,13 @@ export class Binder extends ParseTreeWalker {
                     }
 
                     let submoduleFallback: AliasDeclaration | undefined;
+                    let loadSymbolsFromPath = true;
                     if (implicitImport) {
                         submoduleFallback = {
                             type: DeclarationType.Alias,
                             node: importSymbolNode,
                             path: implicitImport.path,
+                            loadSymbolsFromPath: true,
                             range: getEmptyRange(),
                             usesLocalName: false,
                             moduleName: this._fileInfo.moduleName,
@@ -1669,7 +1674,7 @@ export class Binder extends ParseTreeWalker {
                             node.module.leadingDots === 1 &&
                             node.module.nameParts.length === 0
                         ) {
-                            resolvedPath = '';
+                            loadSymbolsFromPath = false;
                         }
                     }
 
@@ -1677,6 +1682,7 @@ export class Binder extends ParseTreeWalker {
                         type: DeclarationType.Alias,
                         node: importSymbolNode,
                         path: resolvedPath,
+                        loadSymbolsFromPath,
                         usesLocalName: !!importSymbolNode.alias,
                         symbolName: importedName,
                         submoduleFallback,
@@ -2132,7 +2138,12 @@ export class Binder extends ParseTreeWalker {
         const aliasDecl = varSymbol.getDeclarations().find((decl) => decl.type === DeclarationType.Alias) as
             | AliasDeclaration
             | undefined;
-        const resolvedPath = aliasDecl?.path || aliasDecl?.submoduleFallback?.path;
+        const resolvedPath =
+            aliasDecl?.path && aliasDecl.loadSymbolsFromPath
+                ? aliasDecl.path
+                : aliasDecl?.submoduleFallback?.path && aliasDecl.submoduleFallback.loadSymbolsFromPath
+                ? aliasDecl.submoduleFallback.path
+                : undefined;
         if (!resolvedPath) {
             return undefined;
         }
@@ -2182,7 +2193,8 @@ export class Binder extends ParseTreeWalker {
                 newDecl = {
                     type: DeclarationType.Alias,
                     node,
-                    path: '',
+                    path: importInfo.resolvedPaths[importInfo.resolvedPaths.length - 1],
+                    loadSymbolsFromPath: false,
                     moduleName: importInfo.importName,
                     range: getEmptyRange(),
                     firstNamePart: firstNamePartValue,
@@ -2194,6 +2206,7 @@ export class Binder extends ParseTreeWalker {
             // name part we're resolving.
             if (importAlias || node.module.nameParts.length === 1) {
                 newDecl.path = importInfo.resolvedPaths[importInfo.resolvedPaths.length - 1];
+                newDecl.loadSymbolsFromPath = true;
                 this._addImplicitImportsToLoaderActions(importInfo, newDecl);
             } else {
                 // Fill in the remaining name parts.
@@ -2213,7 +2226,8 @@ export class Binder extends ParseTreeWalker {
                     if (!loaderActions) {
                         // Allocate a new loader action.
                         loaderActions = {
-                            path: '',
+                            path: importInfo.resolvedPaths[i],
+                            loadSymbolsFromPath: false,
                             implicitImports: new Map<string, ModuleLoaderActions>(),
                         };
                         if (!curLoaderActions.implicitImports) {
@@ -2226,6 +2240,7 @@ export class Binder extends ParseTreeWalker {
                     // implicit imports as well.
                     if (i === node.module.nameParts.length - 1) {
                         loaderActions.path = importInfo.resolvedPaths[i];
+                        loaderActions.loadSymbolsFromPath = true;
                         this._addImplicitImportsToLoaderActions(importInfo, loaderActions);
                     }
 
@@ -2244,6 +2259,7 @@ export class Binder extends ParseTreeWalker {
                 type: DeclarationType.Alias,
                 node,
                 path: '*** unresolved ***',
+                loadSymbolsFromPath: true,
                 range: getEmptyRange(),
                 usesLocalName: !!importAlias,
                 moduleName: '',
@@ -3622,12 +3638,14 @@ export class Binder extends ParseTreeWalker {
                 : undefined;
             if (existingLoaderAction) {
                 existingLoaderAction.path = implicitImport.path;
+                existingLoaderAction.loadSymbolsFromPath = true;
             } else {
                 if (!loaderActions.implicitImports) {
                     loaderActions.implicitImports = new Map<string, ModuleLoaderActions>();
                 }
                 loaderActions.implicitImports.set(implicitImport.name, {
                     path: implicitImport.path,
+                    loadSymbolsFromPath: true,
                     implicitImports: new Map<string, ModuleLoaderActions>(),
                 });
             }

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

@@ -147,6 +147,9 @@ export interface AliasDeclaration extends DeclarationBase {
     // rename references.
     usesLocalName: boolean;
 
+    // Indicate whether symbols can be loaded from the path.
+    loadSymbolsFromPath: boolean;
+
     // The name of the symbol being imported (used for "from X import Y"
     // statements, not applicable to "import X" statements).
     symbolName?: string | undefined;
@@ -180,6 +183,9 @@ export interface ModuleLoaderActions {
     // a directory).
     path: string;
 
+    // Indicate whether symbols can be loaded from the path.
+    loadSymbolsFromPath: boolean;
+
     // See comment for "implicitImports" field in AliasDeclaration.
     implicitImports?: Map<string, ModuleLoaderActions>;
 }

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

@@ -7,8 +7,9 @@
  * Collection of static methods that operate on declarations.
  */
 
+import { getEmptyRange } from '../common/textRange';
 import { ParseNodeType } from '../parser/parseNodes';
-import { Declaration, DeclarationType, isAliasDeclaration } from './declaration';
+import { AliasDeclaration, Declaration, DeclarationType, isAliasDeclaration, ModuleLoaderActions } from './declaration';
 import { getFileInfoFromNode } from './parseTreeUtils';
 
 export function hasTypeForDeclaration(declaration: Declaration): boolean {
@@ -52,7 +53,11 @@ export function hasTypeForDeclaration(declaration: Declaration): boolean {
     }
 }
 
-export function areDeclarationsSame(decl1: Declaration, decl2: Declaration): boolean {
+export function areDeclarationsSame(
+    decl1: Declaration,
+    decl2: Declaration,
+    treatModuleInImportAndFromImportSame = false
+): boolean {
     if (decl1.type !== decl2.type) {
         return false;
     }
@@ -71,11 +76,22 @@ export function areDeclarationsSame(decl1: Declaration, decl2: Declaration): boo
     // Alias declarations refer to the entire import statement.
     // We need to further differentiate.
     if (decl1.type === DeclarationType.Alias && decl2.type === DeclarationType.Alias) {
-        if (
-            decl1.symbolName !== decl2.symbolName ||
-            decl1.firstNamePart !== decl2.firstNamePart ||
-            decl1.usesLocalName !== decl2.usesLocalName
-        ) {
+        if (decl1.symbolName !== decl2.symbolName || decl1.usesLocalName !== decl2.usesLocalName) {
+            return false;
+        }
+
+        if (treatModuleInImportAndFromImportSame) {
+            // Treat "module" in "import [|module|]", "from [|module|] import ..."
+            // or "from ... import [|module|]" same in IDE services.
+            //
+            // Some case such as "from [|module|] import ...", symbol for [|module|] doesn't even
+            // exist and it can't be referenced inside of a module, but nonetheless, IDE still
+            // needs these sometimes for things like hover tooltip, highlight references,
+            // find all references and etc.
+            return true;
+        }
+
+        if (decl1.firstNamePart !== decl2.firstNamePart) {
             return false;
         }
     }
@@ -162,3 +178,31 @@ export function isDefinedInFile(decl: Declaration, filePath: string) {
     // Other decls, the path points to the file the symbol is defined in.
     return decl.path === filePath;
 }
+
+export function getDeclarationsWithUsesLocalNameRemoved(decls: Declaration[]) {
+    // Make a shallow copy and clear the "usesLocalName" field.
+    return decls.map((localDecl) => {
+        if (localDecl.type !== DeclarationType.Alias) {
+            return localDecl;
+        }
+
+        const nonLocalDecl: AliasDeclaration = { ...localDecl };
+        nonLocalDecl.usesLocalName = false;
+        return nonLocalDecl;
+    });
+}
+
+export function createSynthesizedAliasDeclaration(path: string): AliasDeclaration {
+    // The only time this decl is used is for IDE services such as
+    // the find all references, hover provider and etc.
+    return {
+        type: DeclarationType.Alias,
+        node: undefined!,
+        path,
+        loadSymbolsFromPath: false,
+        range: getEmptyRange(),
+        implicitImports: new Map<string, ModuleLoaderActions>(),
+        usesLocalName: false,
+        moduleName: '',
+    };
+}

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

@@ -28,7 +28,10 @@ import {
     getPathComponents,
     getRelativePathComponentsFromDirectory,
     isDirectory,
+    isDiskPathRoot,
     isFile,
+    normalizePath,
+    normalizePathCase,
     resolvePaths,
     stripFileExtension,
     stripTrailingDirectorySeparator,
@@ -41,6 +44,7 @@ import * as StringUtils from '../common/stringUtils';
 import { isIdentifierChar, isIdentifierStartChar } from '../parser/characters';
 import { PyrightFileSystem } from '../pyrightFileSystem';
 import { ImplicitImport, ImportResult, ImportType } from './importResult';
+import { ImportPath, ParentDirectoryCache } from './parentDirectoryCache';
 import * as PythonPathUtils from './pythonPathUtils';
 import { getPyTypedInfo, PyTypedInfo } from './pyTypedUtils';
 import { isDunderName } from './symbolNameUtils';
@@ -58,6 +62,14 @@ export interface ModuleNameAndType {
     isLocalTypingsFile: boolean;
 }
 
+export function createImportedModuleDescriptor(moduleName: string): ImportedModuleDescriptor {
+    return {
+        leadingDots: 0,
+        nameParts: moduleName.split('.'),
+        importedSymbols: [],
+    };
+}
+
 type CachedImportResults = Map<string, ImportResult>;
 interface SupportedVersionRange {
     min: PythonVersion;
@@ -86,15 +98,21 @@ export class ImportResolver {
     private _cachedTypeshedThirdPartyPackageRoots: string[] | undefined;
     private _cachedEntriesForPath = new Map<string, Dirent[]>();
 
+    protected cachedParentImportResults: ParentDirectoryCache;
+
     constructor(
         public readonly fileSystem: FileSystem,
         protected _configOptions: ConfigOptions,
         public readonly host: Host
-    ) {}
+    ) {
+        this.cachedParentImportResults = new ParentDirectoryCache(() => this.getPythonSearchPaths([]));
+    }
 
     invalidateCache() {
         this._cachedImportResults = new Map<string | undefined, CachedImportResults>();
         this._cachedModuleNameResults = new Map<string, Map<string, ModuleNameAndType>>();
+        this.cachedParentImportResults.reset();
+
         this._invalidateFileSystemCache();
 
         if (this.fileSystem instanceof PyrightFileSystem) {
@@ -123,7 +141,86 @@ export class ImportResolver {
     ): ImportResult {
         const importName = this.formatImportName(moduleDescriptor);
         const importFailureInfo: string[] = [];
+        const importResult = this._resolveImportStrict(
+            importName,
+            sourceFilePath,
+            execEnv,
+            moduleDescriptor,
+            importFailureInfo
+        );
 
+        if (importResult.isImportFound || moduleDescriptor.leadingDots > 0) {
+            return importResult;
+        }
+
+        // If the import is absolute and no other method works, try resolving the
+        // absolute in the importing file's directory, then the parent directory,
+        // and so on, until the import root is reached.
+        sourceFilePath = normalizePathCase(this.fileSystem, normalizePath(sourceFilePath));
+        const origin = ensureTrailingDirectorySeparator(getDirectoryPath(sourceFilePath));
+
+        const result = this.cachedParentImportResults.getImportResult(origin, importName, importResult);
+        if (result) {
+            // Already ran the parent directory resolution for this import name on this location.
+            return this.filterImplicitImports(result, moduleDescriptor.importedSymbols);
+        }
+
+        // Check whether the given file is in the parent directory import resolution cache.
+        const root = this.getParentImportResolutionRoot(sourceFilePath, execEnv.root);
+        if (!this.cachedParentImportResults.checkValidPath(this.fileSystem, sourceFilePath, root)) {
+            return importResult;
+        }
+
+        const importPath: ImportPath = { importPath: undefined };
+
+        // Going up the given folder one by one until we can resolve the import.
+        let current = origin;
+        while (this._shouldWalkUp(current, root, execEnv)) {
+            const result = this.resolveAbsoluteImport(
+                current,
+                execEnv,
+                moduleDescriptor,
+                importName,
+                [],
+                /* allowPartial */ undefined,
+                /* allowNativeLib */ undefined,
+                /* useStubPackage */ false,
+                /* allowPyi */ true
+            );
+
+            this.cachedParentImportResults.checked(current, importName, importPath);
+
+            if (result.isImportFound) {
+                // This will make cache to point to actual path that contains the module we found
+                importPath.importPath = current;
+
+                this.cachedParentImportResults.add({
+                    importResult: result,
+                    path: current,
+                    importName,
+                });
+
+                return this.filterImplicitImports(result, moduleDescriptor.importedSymbols);
+            }
+
+            let success;
+            [success, current] = this._tryWalkUp(current);
+            if (!success) {
+                break;
+            }
+        }
+
+        this.cachedParentImportResults.checked(current, importName, importPath);
+        return importResult;
+    }
+
+    private _resolveImportStrict(
+        importName: string,
+        sourceFilePath: string,
+        execEnv: ExecutionEnvironment,
+        moduleDescriptor: ImportedModuleDescriptor,
+        importFailureInfo: string[]
+    ) {
         const notFoundResult: ImportResult = {
             importName,
             isRelative: false,
@@ -194,6 +291,32 @@ export class ImportResolver {
         sourceFilePath: string,
         execEnv: ExecutionEnvironment,
         moduleDescriptor: ImportedModuleDescriptor
+    ) {
+        const suggestions = this._getCompletionSuggestionsStrict(sourceFilePath, execEnv, moduleDescriptor);
+
+        const root = this.getParentImportResolutionRoot(sourceFilePath, execEnv.root);
+        const origin = ensureTrailingDirectorySeparator(
+            getDirectoryPath(normalizePathCase(this.fileSystem, normalizePath(sourceFilePath)))
+        );
+
+        let current = origin;
+        while (this._shouldWalkUp(current, root, execEnv)) {
+            this.getCompletionSuggestionsAbsolute(current, moduleDescriptor, suggestions, sourceFilePath, execEnv);
+
+            let success;
+            [success, current] = this._tryWalkUp(current);
+            if (!success) {
+                break;
+            }
+        }
+
+        return suggestions;
+    }
+
+    private _getCompletionSuggestionsStrict(
+        sourceFilePath: string,
+        execEnv: ExecutionEnvironment,
+        moduleDescriptor: ImportedModuleDescriptor
     ): Set<string> {
         const importFailureInfo: string[] = [];
         const suggestions = new Set<string>();
@@ -390,6 +513,7 @@ export class ImportResolver {
         // Look for it in the root directory of the execution environment.
         if (execEnv.root) {
             moduleName = this.getModuleNameFromPath(execEnv.root, filePath);
+            importType = ImportType.Local;
         }
 
         for (const extraPath of execEnv.extraPaths) {
@@ -1955,6 +2079,31 @@ export class ImportResolver {
     private _isNativeModuleFileExtension(fileExtension: string): boolean {
         return supportedNativeLibExtensions.some((ext) => ext === fileExtension);
     }
+
+    private _tryWalkUp(current: string): [success: boolean, path: string] {
+        if (isDiskPathRoot(current)) {
+            return [false, ''];
+        }
+
+        return [
+            true,
+            ensureTrailingDirectorySeparator(
+                normalizePathCase(this.fileSystem, normalizePath(combinePaths(current, '..')))
+            ),
+        ];
+    }
+
+    private _shouldWalkUp(current: string, root: string, execEnv: ExecutionEnvironment) {
+        return current.length > root.length || (current === root && !execEnv.root);
+    }
+
+    protected getParentImportResolutionRoot(sourceFilePath: string, executionRoot: string | undefined) {
+        if (executionRoot) {
+            return ensureTrailingDirectorySeparator(normalizePathCase(this.fileSystem, normalizePath(executionRoot)));
+        }
+
+        return ensureTrailingDirectorySeparator(getDirectoryPath(sourceFilePath));
+    }
 }
 
 export type ImportResolverFactory = (fs: FileSystem, options: ConfigOptions, host: Host) => ImportResolver;

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

@@ -154,7 +154,11 @@ export function getTextEditsForAutoImportSymbolAddition(
     parseResults: ParseResults
 ): TextEditAction[] {
     const additionEdits: AdditionEdit[] = [];
-    if (!importStatement.node || importStatement.node.nodeType !== ParseNodeType.ImportFrom) {
+    if (
+        !importStatement.node ||
+        importStatement.node.nodeType !== ParseNodeType.ImportFrom ||
+        importStatement.node.isWildcardImport
+    ) {
         return additionEdits;
     }
 

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

@@ -16,7 +16,7 @@ import { FullAccessHost } from '../common/fullAccessHost';
 import { combinePaths, getDirectoryPath, getFileExtension, stripFileExtension, tryStat } from '../common/pathUtils';
 import { getEmptyRange, Range } from '../common/textRange';
 import { DeclarationType, FunctionDeclaration, VariableDeclaration } from './declaration';
-import { ImportedModuleDescriptor, ImportResolver } from './importResolver';
+import { createImportedModuleDescriptor, ImportResolver } from './importResolver';
 import {
     AlternateSymbolNameMap,
     getEmptyReport,
@@ -180,12 +180,7 @@ export class PackageTypeVerifier {
     }
 
     private _resolveImport(moduleName: string) {
-        const moduleDescriptor: ImportedModuleDescriptor = {
-            leadingDots: 0,
-            nameParts: moduleName.split('.'),
-            importedSymbols: [],
-        };
-        return this._importResolver.resolveImport('', this._execEnv, moduleDescriptor);
+        return this._importResolver.resolveImport('', this._execEnv, createImportedModuleDescriptor(moduleName));
     }
 
     private _getPublicSymbolsForModule(
@@ -1147,13 +1142,11 @@ export class PackageTypeVerifier {
     }
 
     private _getDirectoryForPackage(packageName: string): string | undefined {
-        const moduleDescriptor: ImportedModuleDescriptor = {
-            leadingDots: 0,
-            nameParts: [packageName],
-            importedSymbols: [],
-        };
-
-        const importResult = this._importResolver.resolveImport('', this._execEnv, moduleDescriptor);
+        const importResult = this._importResolver.resolveImport(
+            '',
+            this._execEnv,
+            createImportedModuleDescriptor(packageName)
+        );
 
         if (importResult.isImportFound) {
             const resolvedPath = importResult.resolvedPaths[importResult.resolvedPaths.length - 1];

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

@@ -0,0 +1,88 @@
+/*
+ * parentDirectoryCache.ts
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ *
+ * Cache to hold parent directory import result to make sure
+ * we don't repeatedly search folders.
+ */
+
+import { getOrAdd } from '../common/collectionUtils';
+import { FileSystem } from '../common/fileSystem';
+import { ensureTrailingDirectorySeparator, normalizePath, normalizePathCase } from '../common/pathUtils';
+import { ImportResult } from './importResult';
+
+export type ImportPath = { importPath: string | undefined };
+
+type CacheEntry = { importResult: ImportResult; path: string; importName: string };
+
+export class ParentDirectoryCache {
+    private readonly _importChecked = new Map<string, Map<string, ImportPath>>();
+    private readonly _cachedResults = new Map<string, Map<string, ImportResult>>();
+
+    private _libPathCache: string[] | undefined = undefined;
+
+    constructor(private _importRootGetter: () => string[]) {
+        // empty
+    }
+
+    getImportResult(path: string, importName: string, importResult: ImportResult): ImportResult | undefined {
+        const result = this._cachedResults.get(importName)?.get(path);
+        if (result) {
+            // We already checked for the importName at the path.
+            // Return the result if succeeded otherwise, return regular import result given.
+            return result ?? importResult;
+        }
+
+        const checked = this._importChecked.get(importName)?.get(path);
+        if (checked) {
+            // We already checked for the importName at the path.
+            if (!checked.importPath) {
+                return importResult;
+            }
+
+            return this._cachedResults.get(importName)?.get(checked.importPath) ?? importResult;
+        }
+
+        return undefined;
+    }
+
+    checkValidPath(fs: FileSystem, sourceFilePath: string, root: string): boolean {
+        if (!sourceFilePath.startsWith(root)) {
+            // We don't search containing folders for libs.
+            return false;
+        }
+
+        this._libPathCache =
+            this._libPathCache ??
+            this._importRootGetter()
+                .map((r) => ensureTrailingDirectorySeparator(normalizePathCase(fs, normalizePath(r))))
+                .filter((r) => r !== root)
+                .filter((r) => r.startsWith(root));
+
+        if (this._libPathCache.some((p) => sourceFilePath.startsWith(p))) {
+            // Make sure it is not lib folders under user code root.
+            // ex) .venv folder
+            return false;
+        }
+
+        return true;
+    }
+
+    checked(path: string, importName: string, importPath: ImportPath) {
+        getOrAdd(this._importChecked, importName, () => new Map<string, ImportPath>()).set(path, importPath);
+    }
+
+    add(result: CacheEntry) {
+        getOrAdd(this._cachedResults, result.importName, () => new Map<string, ImportResult>()).set(
+            result.path,
+            result.importResult
+        );
+    }
+
+    reset() {
+        this._importChecked.clear();
+        this._cachedResults.clear();
+        this._libPathCache = undefined;
+    }
+}

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

@@ -19,6 +19,7 @@ import {
 } from 'vscode-languageserver-types';
 
 import { OperationCanceledException, throwIfCancellationRequested } from '../common/cancellationUtils';
+import { removeArrayElements } from '../common/collectionUtils';
 import { ConfigOptions, ExecutionEnvironment } from '../common/configOptions';
 import { ConsoleInterface, StandardConsole } from '../common/console';
 import { assert } from '../common/debug';
@@ -53,11 +54,14 @@ import { DefinitionFilter } from '../languageService/definitionProvider';
 import { IndexOptions, IndexResults, WorkspaceSymbolCallback } from '../languageService/documentSymbolProvider';
 import { HoverResults } from '../languageService/hoverProvider';
 import { ReferenceCallback, ReferencesResult } from '../languageService/referencesProvider';
+import { RenameModuleProvider } from '../languageService/renameModuleProvider';
 import { SignatureHelpResults } from '../languageService/signatureHelpProvider';
+import { ParseNodeType } from '../parser/parseNodes';
 import { ParseResults } from '../parser/parser';
 import { ImportLookupResult } from './analyzerFileInfo';
 import * as AnalyzerNodeInfo from './analyzerNodeInfo';
 import { CircularDependency } from './circularDependency';
+import { isAliasDeclaration } from './declaration';
 import { ImportResolver } from './importResolver';
 import { ImportResult, ImportType } from './importResult';
 import { findNodeByOffset, getDocString } from './parseTreeUtils';
@@ -374,6 +378,10 @@ export class Program {
         return this._sourceFileList.filter((s) => s.isTracked);
     }
 
+    getOpened(): SourceFileInfo[] {
+        return this._sourceFileList.filter((s) => s.isOpenByClient);
+    }
+
     getFilesToAnalyzeCount() {
         let sourceFileCount = 0;
 
@@ -1027,7 +1035,7 @@ export class Program {
         }
 
         const sourceFile = sourceFileInfo.sourceFile;
-        const fileContents = sourceFile.getFileContents();
+        const fileContents = sourceFile.getOpenFileContents();
         if (fileContents === undefined) {
             // this only works with opened file
             return undefined;
@@ -1062,7 +1070,7 @@ export class Program {
         }
 
         const sourceFile = sourceFileInfo.sourceFile;
-        const fileContents = sourceFile.getFileContents();
+        const fileContents = sourceFile.getOpenFileContents();
         if (fileContents === undefined) {
             // this only works with opened file
             return [];
@@ -1319,22 +1327,16 @@ export class Program {
                 return undefined;
             }
 
-            let content: string | undefined = undefined;
+            const content = sourceFileInfo.sourceFile.getFileContent() ?? '';
             if (
                 options.indexingForAutoImportMode &&
                 !sourceFileInfo.sourceFile.isStubFile() &&
-                !sourceFileInfo.sourceFile.isThirdPartyPyTypedPresent() &&
-                sourceFileInfo.sourceFile.getClientVersion() === undefined
+                !sourceFileInfo.sourceFile.isThirdPartyPyTypedPresent()
             ) {
-                try {
-                    // Perf optimization. if py file doesn't contain __all__
-                    // No need to parse and bind.
-                    content = this._fs.readFileSync(filePath, 'utf8');
-                    if (content.indexOf('__all__') < 0) {
-                        return undefined;
-                    }
-                } catch (error) {
-                    content = undefined;
+                // Perf optimization. if py file doesn't contain __all__
+                // No need to parse and bind.
+                if (content.indexOf('__all__') < 0) {
+                    return undefined;
                 }
             }
 
@@ -1568,6 +1570,57 @@ export class Program {
         });
     }
 
+    renameModule(filePath: string, newFilePath: string, token: CancellationToken): FileEditAction[] | undefined {
+        return this._runEvaluatorWithCancellationToken(token, () => {
+            const fileInfo = this._getSourceFileInfoFromPath(filePath);
+            if (!fileInfo) {
+                return undefined;
+            }
+
+            const renameModuleProvider = RenameModuleProvider.create(
+                this._importResolver,
+                this._configOptions,
+                this._evaluator!,
+                filePath,
+                newFilePath,
+                token
+            );
+            if (!renameModuleProvider) {
+                return undefined;
+            }
+
+            // _sourceFileList contains every user files that match "include" pattern including
+            // py file even if corresponding pyi exists.
+            for (const currentFileInfo of this._sourceFileList) {
+                // Make sure we only touch user code to prevent us
+                // from accidentally changing third party library or type stub.
+                if (!this._isUserCode(currentFileInfo)) {
+                    continue;
+                }
+
+                // If module name isn't mentioned in the current file, skip the file.
+                const content = currentFileInfo.sourceFile.getFileContent() ?? '';
+                if (content.indexOf(renameModuleProvider.symbolName) < 0) {
+                    continue;
+                }
+
+                this._bindFile(currentFileInfo, content);
+                const parseResult = currentFileInfo.sourceFile.getParseResults();
+                if (!parseResult) {
+                    continue;
+                }
+
+                renameModuleProvider.renameModuleReferences(currentFileInfo.sourceFile.getFilePath(), parseResult);
+
+                // This operation can consume significant memory, so check
+                // for situations where we need to discard the type cache.
+                this._handleMemoryHighUsage();
+            }
+
+            return renameModuleProvider.getEdits();
+        });
+    }
+
     renameSymbolAtPosition(
         filePath: string,
         position: Position,
@@ -1596,6 +1649,36 @@ export class Program {
                 return undefined;
             }
 
+            // We only allow renaming module alias, filter out any other alias decls.
+            removeArrayElements(referencesResult.declarations, (d) => {
+                if (!isAliasDeclaration(d)) {
+                    return false;
+                }
+
+                // We must have alias and decl node that point to import statement.
+                if (!d.usesLocalName || !d.node) {
+                    return true;
+                }
+
+                // d.node can't be ImportFrom if usesLocalName is true.
+                // but we are doing this for type checker.
+                if (d.node.nodeType === ParseNodeType.ImportFrom) {
+                    return true;
+                }
+
+                // Check alias and what we are renaming is same thing.
+                if (d.node.alias?.value !== referencesResult.symbolName) {
+                    return true;
+                }
+
+                return false;
+            });
+
+            if (referencesResult.declarations.length === 0) {
+                // There is no symbol we can rename.
+                return undefined;
+            }
+
             if (
                 !isDefaultWorkspace &&
                 referencesResult.declarations.some((d) => !this._isUserCode(this._getSourceFileInfoFromPath(d.path)))
@@ -1604,11 +1687,6 @@ export class Program {
                 return undefined;
             }
 
-            if (referencesResult.declarations.length === 0) {
-                // There is no symbol we can rename.
-                return undefined;
-            }
-
             // Do we need to do a global search as well?
             if (referencesResult.requiresGlobalSearch && !isDefaultWorkspace) {
                 for (const curSourceFileInfo of this._sourceFileList) {
@@ -1783,6 +1861,10 @@ export class Program {
         return sourceFileInfo.sourceFile.performQuickAction(command, args, token);
     }
 
+    test_createSourceMapper(execEnv: ExecutionEnvironment) {
+        return this._createSourceMapper(execEnv, /*mapCompiled*/ false);
+    }
+
     private _handleMemoryHighUsage() {
         const typeCacheSize = this._evaluator!.getTypeCacheSize();
 

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

@@ -60,7 +60,7 @@ import { ReferenceCallback } from '../languageService/referencesProvider';
 import { SignatureHelpResults } from '../languageService/signatureHelpProvider';
 import { AnalysisCompleteCallback } from './analysis';
 import { BackgroundAnalysisProgram, BackgroundAnalysisProgramFactory } from './backgroundAnalysisProgram';
-import { ImportedModuleDescriptor, ImportResolver, ImportResolverFactory } from './importResolver';
+import { createImportedModuleDescriptor, ImportResolver, ImportResolverFactory } from './importResolver';
 import { MaxAnalysisTime } from './program';
 import { findPythonSearchPaths } from './pythonPathUtils';
 import { TypeEvaluator } from './typeEvaluatorTypes';
@@ -148,10 +148,10 @@ export class AnalyzerService {
                   );
     }
 
-    clone(instanceName: string, backgroundAnalysis?: BackgroundAnalysisBase): AnalyzerService {
-        return new AnalyzerService(
+    clone(instanceName: string, backgroundAnalysis?: BackgroundAnalysisBase, fs?: FileSystem): AnalyzerService {
+        const service = new AnalyzerService(
             instanceName,
-            this._fs,
+            fs ?? this._fs,
             this._console,
             this._hostFactory,
             this._importResolverFactory,
@@ -162,6 +162,20 @@ export class AnalyzerService {
             this._backgroundAnalysisProgramFactory,
             this._cancellationProvider
         );
+
+        // Make sure we keep editor content (open file) which could be different than one in the file system.
+        for (const fileInfo of this.backgroundAnalysisProgram.program.getOpened()) {
+            const version = fileInfo.sourceFile.getClientVersion();
+            if (version !== undefined) {
+                service.setFileOpened(
+                    fileInfo.sourceFile.getFilePath(),
+                    version,
+                    fileInfo.sourceFile.getOpenFileContents()!
+                );
+            }
+        }
+
+        return service;
     }
 
     dispose() {
@@ -187,7 +201,7 @@ export class AnalyzerService {
         this._backgroundAnalysisProgram.setCompletionCallback(callback);
     }
 
-    setOptions(commandLineOptions: CommandLineOptions, reanalyze = true): void {
+    setOptions(commandLineOptions: CommandLineOptions): void {
         this._commandLineOptions = commandLineOptions;
 
         const host = this._hostFactory();
@@ -205,7 +219,7 @@ export class AnalyzerService {
         this._executionRootPath = normalizePath(
             combinePaths(commandLineOptions.executionRoot, configOptions.projectRoot)
         );
-        this._applyConfigOptions(host, reanalyze);
+        this._applyConfigOptions(host);
     }
 
     setFileOpened(path: string, version: number | null, contents: string) {
@@ -736,7 +750,11 @@ export class AnalyzerService {
     // This is called after a new type stub has been created. It allows
     // us to invalidate caches and force reanalysis of files that potentially
     // are affected by the appearance of a new type stub.
-    invalidateAndForceReanalysis(rebuildLibraryIndexing = true) {
+    invalidateAndForceReanalysis(rebuildLibraryIndexing = true, updateTrackedFileList = false) {
+        if (updateTrackedFileList) {
+            this._updateTrackedFileList(/* markFilesDirtyUnconditionally */ false);
+        }
+
         // Mark all files with one or more errors dirty.
         this._backgroundAnalysisProgram.invalidateAndForceReanalysis(rebuildLibraryIndexing);
     }
@@ -948,12 +966,7 @@ export class AnalyzerService {
         // for a different set of files.
         if (this._typeStubTargetImportName) {
             const execEnv = this._configOptions.findExecEnvironment(this._executionRootPath);
-            const moduleDescriptor: ImportedModuleDescriptor = {
-                leadingDots: 0,
-                nameParts: this._typeStubTargetImportName.split('.'),
-                importedSymbols: [],
-            };
-
+            const moduleDescriptor = createImportedModuleDescriptor(this._typeStubTargetImportName);
             const importResult = this._backgroundAnalysisProgram.importResolver.resolveImport(
                 '',
                 execEnv,
@@ -1369,7 +1382,7 @@ export class AnalyzerService {
         }
     }
 
-    private _applyConfigOptions(host: Host, reanalyze = true) {
+    private _applyConfigOptions(host: Host) {
         // Allocate a new import resolver because the old one has information
         // cached based on the previous config options.
         const importResolver = this._importResolverFactory(
@@ -1396,9 +1409,7 @@ export class AnalyzerService {
         this._updateSourceFileWatchers();
         this._updateTrackedFileList(true);
 
-        if (reanalyze) {
-            this._scheduleReanalysis(false);
-        }
+        this._scheduleReanalysis(false);
     }
 
     private _clearReanalysisTimer() {
@@ -1409,7 +1420,7 @@ export class AnalyzerService {
     }
 
     private _scheduleReanalysis(requireTrackedFileUpdate: boolean) {
-        if (this._disposed) {
+        if (this._disposed || !this._commandLineOptions?.enableAmbientAnalysis) {
             // already disposed
             return;
         }

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

@@ -43,7 +43,7 @@ import { performQuickAction } from '../languageService/quickActions';
 import { ReferenceCallback, ReferencesProvider, ReferencesResult } from '../languageService/referencesProvider';
 import { SignatureHelpProvider, SignatureHelpResults } from '../languageService/signatureHelpProvider';
 import { Localizer } from '../localization/localize';
-import { ModuleNode } from '../parser/parseNodes';
+import { ModuleNode, NameNode } from '../parser/parseNodes';
 import { ModuleImport, ParseOptions, Parser, ParseResults } from '../parser/parser';
 import { Token } from '../parser/tokenizerTypes';
 import { AnalyzerFileInfo, ImportLookup } from './analyzerFileInfo';
@@ -414,10 +414,35 @@ export class SourceFile {
         return this._clientDocument?.version;
     }
 
-    getFileContents() {
+    getOpenFileContents() {
         return this._clientDocument?.getText();
     }
 
+    getFileContent(): string | undefined {
+        // Get current buffer content if the file is opened.
+        const openFileContent = this.getOpenFileContents();
+        if (openFileContent) {
+            return openFileContent;
+        }
+
+        // Otherwise, get content from file system.
+        try {
+            // Check the file's length before attempting to read its full contents.
+            const fileStat = this.fileSystem.statSync(this._filePath);
+            if (fileStat.size > _maxSourceFileSize) {
+                this._console.error(
+                    `File length of "${this._filePath}" is ${fileStat.size} ` +
+                        `which exceeds the maximum supported file size of ${_maxSourceFileSize}`
+                );
+                throw new Error('File larger than max');
+            }
+
+            return this.fileSystem.readFileSync(this._filePath, 'utf8');
+        } catch (error) {
+            return undefined;
+        }
+    }
+
     setClientVersion(version: number | null, contents: TextDocumentContentChangeEvent[]): void {
         if (version === null) {
             this._clientDocument = undefined;
@@ -522,23 +547,16 @@ export class SourceFile {
             }
 
             const diagSink = new DiagnosticSink();
-            let fileContents = this.getFileContents();
+            let fileContents = this.getOpenFileContents();
             if (fileContents === undefined) {
                 try {
                     const startTime = timingStats.readFileTime.totalTime;
                     timingStats.readFileTime.timeOperation(() => {
-                        // Check the file's length before attempting to read its full contents.
-                        const fileStat = this.fileSystem.statSync(this._filePath);
-                        if (fileStat.size > _maxSourceFileSize) {
-                            this._console.error(
-                                `File length of "${this._filePath}" is ${fileStat.size} ` +
-                                    `which exceeds the maximum supported file size of ${_maxSourceFileSize}`
-                            );
-                            throw new Error('File larger than max');
-                        }
-
                         // Read the file's contents.
-                        fileContents = content ?? this.fileSystem.readFileSync(this._filePath, 'utf8');
+                        fileContents = content ?? this.getFileContent();
+                        if (!fileContents) {
+                            throw new Error("Can't get file content");
+                        }
 
                         // Remember the length and hash for comparison purposes.
                         this._lastFileContentLength = fileContents.length;
@@ -700,6 +718,21 @@ export class SourceFile {
         );
     }
 
+    getDeclarationForNode(
+        sourceMapper: SourceMapper,
+        node: NameNode,
+        evaluator: TypeEvaluator,
+        reporter: ReferenceCallback | undefined,
+        token: CancellationToken
+    ): ReferencesResult | undefined {
+        // If we have no completed analysis job, there's nothing to do.
+        if (!this._parseResults) {
+            return undefined;
+        }
+
+        return ReferencesProvider.getDeclarationForNode(sourceMapper, this._filePath, node, evaluator, reporter, token);
+    }
+
     getDeclarationForPosition(
         sourceMapper: SourceMapper,
         position: Position,
@@ -847,7 +880,7 @@ export class SourceFile {
 
         // This command should be called only for open files, in which
         // case we should have the file contents already loaded.
-        const fileContents = this.getFileContents();
+        const fileContents = this.getOpenFileContents();
         if (fileContents === undefined) {
             return undefined;
         }
@@ -888,7 +921,7 @@ export class SourceFile {
         completionItem: CompletionItem,
         token: CancellationToken
     ) {
-        const fileContents = this.getFileContents();
+        const fileContents = this.getOpenFileContents();
         if (!this._parseResults || fileContents === undefined) {
             return;
         }

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

@@ -24,7 +24,7 @@ import { AddMissingOptionalToParamAction, DiagnosticAddendum } from '../common/d
 import { DiagnosticRule } from '../common/diagnosticRules';
 import { convertOffsetsToRange } from '../common/positionUtils';
 import { PythonVersion } from '../common/pythonVersion';
-import { getEmptyRange, TextRange } from '../common/textRange';
+import { TextRange } from '../common/textRange';
 import { Localizer } from '../localization/localize';
 import {
     ArgumentCategory,
@@ -103,7 +103,6 @@ import {
     validateDataClassTransformDecorator,
 } from './dataClasses';
 import {
-    AliasDeclaration,
     ClassDeclaration,
     Declaration,
     DeclarationType,
@@ -111,7 +110,12 @@ import {
     ModuleLoaderActions,
     VariableDeclaration,
 } from './declaration';
-import { isExplicitTypeAliasDeclaration, isPossibleTypeAliasDeclaration } from './declarationUtils';
+import {
+    createSynthesizedAliasDeclaration,
+    getDeclarationsWithUsesLocalNameRemoved,
+    isExplicitTypeAliasDeclaration,
+    isPossibleTypeAliasDeclaration,
+} from './declarationUtils';
 import { createNamedTupleType } from './namedTuples';
 import * as ParseTreeUtils from './parseTreeUtils';
 import { assignTypeToPatternTargets, narrowTypeBasedOnPattern } from './patternMatching';
@@ -16266,17 +16270,7 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
                         return decl.type === DeclarationType.Alias && decl.node === node.parent;
                     });
 
-                    // Make a shallow copy and clear the "usesLocalName" field.
-                    const nonLocalDecls = declsForThisImport.map((localDecl) => {
-                        if (localDecl.type === DeclarationType.Alias) {
-                            const nonLocalDecl: AliasDeclaration = { ...localDecl };
-                            nonLocalDecl.usesLocalName = false;
-                            return nonLocalDecl;
-                        }
-                        return localDecl;
-                    });
-
-                    declarations.push(...nonLocalDecls);
+                    declarations.push(...getDeclarationsWithUsesLocalNameRemoved(declsForThisImport));
                 }
             }
         } else if (
@@ -16351,17 +16345,9 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
                     evaluateTypesForStatement(node);
 
                     // Synthesize an alias declaration for this name part. The only
-                    // time this case is used is for the hover provider.
-                    const aliasDeclaration: AliasDeclaration = {
-                        type: DeclarationType.Alias,
-                        node: undefined!,
-                        path: importInfo.resolvedPaths[namePartIndex],
-                        range: getEmptyRange(),
-                        implicitImports: new Map<string, ModuleLoaderActions>(),
-                        usesLocalName: false,
-                        moduleName: '',
-                    };
-                    declarations.push(aliasDeclaration);
+                    // time this case is used is for IDE services such as
+                    // the find all references, hover provider and etc.
+                    declarations.push(createSynthesizedAliasDeclaration(importInfo.resolvedPaths[namePartIndex]));
                 }
             }
         } else if (node.parent && node.parent.nodeType === ParseNodeType.Argument && node === node.parent.name) {
@@ -16595,7 +16581,7 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
             loaderActions: ModuleLoaderActions,
             importLookup: ImportLookup
         ): Type {
-            if (loaderActions.path) {
+            if (loaderActions.path && loaderActions.loadSymbolsFromPath) {
                 const lookupResults = importLookup(loaderActions.path);
                 if (lookupResults) {
                     moduleType.fields = lookupResults.symbolTable;

packages/pyright-internal/src/commands/createTypeStub.ts

@@ -8,11 +8,8 @@
 
 import { CancellationToken, ExecuteCommandParams } from 'vscode-languageserver';
 
-import { AnalyzerService } from '../analyzer/service';
 import { OperationCanceledException } from '../common/cancellationUtils';
-import { createDeferred } from '../common/deferred';
-import { convertPathToUri } from '../common/pathUtils';
-import { LanguageServerInterface, WorkspaceServiceInstance } from '../languageServerBase';
+import { LanguageServerInterface } from '../languageServerBase';
 import { AnalyzerServiceExecutor } from '../languageService/analyzerServiceExecutor';
 import { ServerCommand } from './commandController';
 
@@ -25,28 +22,20 @@ export class CreateTypeStubCommand implements ServerCommand {
             const importName = cmdParams.arguments[1];
             const callingFile = cmdParams.arguments[2];
 
-            const service = await this._createTypeStubService(callingFile);
-
-            // Allocate a temporary pseudo-workspace to perform this job.
-            const workspace: WorkspaceServiceInstance = {
-                workspaceName: `Create Type Stub ${importName}`,
-                rootPath: workspaceRoot,
-                rootUri: convertPathToUri(this._ls.fs, workspaceRoot),
-                serviceInstance: service,
-                disableLanguageServices: true,
-                disableOrganizeImports: true,
-                isInitialized: createDeferred<boolean>(),
-            };
-
-            const serverSettings = await this._ls.getSettings(workspace);
-            AnalyzerServiceExecutor.runWithOptions(this._ls.rootPath, workspace, serverSettings, importName, false);
+            const service = await AnalyzerServiceExecutor.cloneService(
+                this._ls,
+                await this._ls.getWorkspaceForFile(callingFile ?? workspaceRoot),
+                importName,
+                this._ls.createBackgroundAnalysis()
+            );
 
             try {
                 await service.writeTypeStubInBackground(token);
                 service.dispose();
+
                 const infoMessage = `Type stub was successfully created for '${importName}'.`;
                 this._ls.window.showInformationMessage(infoMessage);
-                this._handlePostCreateTypeStub();
+                this._ls.reanalyze();
             } catch (err) {
                 const isCancellation = OperationCanceledException.is(err);
                 if (isCancellation) {
@@ -64,24 +53,4 @@ export class CreateTypeStubCommand implements ServerCommand {
             }
         }
     }
-
-    // Creates a service instance that's used for creating type
-    // stubs for a specified target library.
-    private async _createTypeStubService(callingFile?: string): Promise<AnalyzerService> {
-        if (callingFile) {
-            // this should let us to inherit all execution env of the calling file
-            // if it is invoked from IDE through code action
-            const workspace = await this._ls.getWorkspaceForFile(callingFile);
-
-            // new service has its own background analysis running on its own thread
-            // to not block main bg running background analysis
-            return workspace.serviceInstance.clone('Type stub', this._ls.createBackgroundAnalysis());
-        }
-
-        return new AnalyzerService('Type stub', this._ls.fs, this._ls.console);
-    }
-
-    private _handlePostCreateTypeStub() {
-        this._ls.reanalyze();
-    }
 }

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

@@ -123,4 +123,7 @@ export class CommandLineOptions {
 
     // Minimum threshold for type eval logging
     typeEvaluationTimeThreshold = 50;
+
+    // Run ambient analysis
+    enableAmbientAnalysis = true;
 }

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

@@ -679,6 +679,10 @@ export function getWildcardRoot(rootPath: string, fileSpec: string): string {
         pathComponents[0] = stripTrailingDirectorySeparator(pathComponents[0]);
     }
 
+    if (pathComponents.length === 1 && !pathComponents[0]) {
+        return path.sep;
+    }
+
     let wildcardRoot = '';
     let firstComponent = true;
 

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

@@ -154,3 +154,8 @@ export function getCharacterCount(value: string, ch: string) {
     }
     return result;
 }
+
+export function getLastDottedString(text: string) {
+    const index = text.lastIndexOf('.');
+    return index > 0 ? text.substring(index + 1) : text;
+}

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

@@ -8,9 +8,13 @@
  * with a specified set of options.
  */
 import { isPythonBinary } from '../analyzer/pythonPathUtils';
+import { AnalyzerService } from '../analyzer/service';
+import type { BackgroundAnalysis } from '../backgroundAnalysis';
 import { CommandLineOptions } from '../common/commandLineOptions';
+import { createDeferred } from '../common/deferred';
+import { FileSystem } from '../common/fileSystem';
 import { combinePaths } from '../common/pathUtils';
-import { ServerSettings, WorkspaceServiceInstance } from '../languageServerBase';
+import { LanguageServerInterface, ServerSettings, WorkspaceServiceInstance } from '../languageServerBase';
 
 export class AnalyzerServiceExecutor {
     static runWithOptions(
@@ -29,7 +33,37 @@ export class AnalyzerServiceExecutor {
         );
 
         // Setting options causes the analyzer service to re-analyze everything.
-        workspace.serviceInstance.setOptions(commandLineOptions, trackFiles);
+        workspace.serviceInstance.setOptions(commandLineOptions);
+    }
+
+    static async cloneService(
+        ls: LanguageServerInterface,
+        workspace: WorkspaceServiceInstance,
+        typeStubTargetImportName?: string,
+        backgroundAnalysis?: BackgroundAnalysis,
+        fileSystem?: FileSystem
+    ): Promise<AnalyzerService> {
+        // Allocate a temporary pseudo-workspace to perform this job.
+        const tempWorkspace: WorkspaceServiceInstance = {
+            workspaceName: `temp workspace for cloned service`,
+            rootPath: workspace.rootPath,
+            rootUri: workspace.rootUri,
+            serviceInstance: workspace.serviceInstance.clone('cloned service', backgroundAnalysis, fileSystem),
+            disableLanguageServices: true,
+            disableOrganizeImports: true,
+            isInitialized: createDeferred<boolean>(),
+        };
+
+        const serverSettings = await ls.getSettings(workspace);
+        AnalyzerServiceExecutor.runWithOptions(
+            ls.rootPath,
+            tempWorkspace,
+            serverSettings,
+            typeStubTargetImportName,
+            /* trackFiles */ false
+        );
+
+        return tempWorkspace.serviceInstance;
     }
 }
 
@@ -48,6 +82,7 @@ function getEffectiveCommandLineOptions(
     commandLineOptions.indexing = serverSettings.indexing;
     commandLineOptions.logTypeEvaluationTime = serverSettings.logTypeEvaluationTime ?? false;
     commandLineOptions.typeEvaluationTimeThreshold = serverSettings.typeEvaluationTimeThreshold ?? 50;
+    commandLineOptions.enableAmbientAnalysis = trackFiles;
 
     if (!trackFiles) {
         commandLineOptions.watchForSourceChanges = false;

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

@@ -711,7 +711,11 @@ export class AutoImporter {
             }
 
             // Does an 'import from' statement already exist?
-            if (importName && importStatement.node.nodeType === ParseNodeType.ImportFrom) {
+            if (
+                importName &&
+                importStatement.node.nodeType === ParseNodeType.ImportFrom &&
+                !importStatement.node.isWildcardImport
+            ) {
                 // If so, see whether what we want already exist.
                 const importNode = importStatement.node.imports.find((i) => i.name.value === importName);
                 if (importNode) {
@@ -744,7 +748,7 @@ export class AutoImporter {
             // If it is the module itself that got imported, make sure we don't import it again.
             // ex) from module import submodule
             const imported = this._importStatements.orderedImports.find((i) => i.moduleName === moduleName);
-            if (imported && imported.node.nodeType === ParseNodeType.ImportFrom) {
+            if (imported && imported.node.nodeType === ParseNodeType.ImportFrom && !imported.node.isWildcardImport) {
                 const importFrom = imported.node.imports.find((i) => i.name.value === importName);
                 if (importFrom) {
                     // For now, we don't check whether alias or moduleName got overwritten at

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

@@ -11,6 +11,7 @@ import { CancellationToken, CodeAction, CodeActionKind, Command } from 'vscode-l
 import { Commands } from '../commands/commands';
 import { throwIfCancellationRequested } from '../common/cancellationUtils';
 import { AddMissingOptionalToParamAction, CreateTypeStubFileAction } from '../common/diagnostic';
+import { convertPathToUri } from '../common/pathUtils';
 import { Range } from '../common/textRange';
 import { WorkspaceServiceInstance } from '../languageServerBase';
 import { Localizer } from '../localization/localize';
@@ -63,11 +64,13 @@ export class CodeActionProvider {
                     .getActions()!
                     .find((a) => a.action === Commands.addMissingOptionalToParam) as AddMissingOptionalToParamAction;
                 if (action) {
+                    const fs = workspace.serviceInstance.getImportResolver().fileSystem;
                     const addMissingOptionalAction = CodeAction.create(
                         Localizer.CodeAction.addOptionalToAnnotation(),
                         Command.create(
                             Localizer.CodeAction.addOptionalToAnnotation(),
                             Commands.addMissingOptionalToParam,
+                            convertPathToUri(fs, filePath),
                             action.offsetOfTypeNode
                         ),
                         CodeActionKind.QuickFix

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

@@ -1559,19 +1559,17 @@ export class CompletionProvider {
             }
         }
 
-        const results: NameNode[] = [];
-        const collector = new DocumentSymbolCollector(
+        const results = DocumentSymbolCollector.collectFromNode(
             indexNode.baseExpression,
             this._evaluator,
-            results,
             this._cancellationToken,
             startingNode
         );
-        collector.collect();
 
         const keys: Set<string> = new Set<string>();
-        for (const nameNode of results) {
-            const node = nameNode.parent?.nodeType === ParseNodeType.TypeAnnotation ? nameNode.parent : nameNode;
+        for (const result of results) {
+            const node =
+                result.node.parent?.nodeType === ParseNodeType.TypeAnnotation ? result.node.parent : result.node;
 
             if (
                 node.parent?.nodeType === ParseNodeType.Assignment ||

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

@@ -15,7 +15,7 @@ import { TypeEvaluator } from '../analyzer/typeEvaluatorTypes';
 import { throwIfCancellationRequested } from '../common/cancellationUtils';
 import { convertOffsetsToRange, convertPositionToOffset } from '../common/positionUtils';
 import { Position, TextRange } from '../common/textRange';
-import { NameNode, ParseNodeType } from '../parser/parseNodes';
+import { ParseNodeType } from '../parser/parseNodes';
 import { ParseResults } from '../parser/parser';
 import { DocumentSymbolCollector } from './documentSymbolCollector';
 
@@ -42,14 +42,20 @@ export class DocumentHighlightProvider {
             return undefined;
         }
 
-        const results: NameNode[] = [];
-        const collector = new DocumentSymbolCollector(node, evaluator, results, token);
+        const results = DocumentSymbolCollector.collectFromNode(
+            node,
+            evaluator,
+            token,
+            parseResults.parseTree,
+            /* treatModuleInImportAndFromImportSame */ true
+        );
 
-        collector.collect();
-
-        return results.map((n) => ({
-            kind: ParseTreeUtils.isWriteAccess(n) ? DocumentHighlightKind.Write : DocumentHighlightKind.Read,
-            range: convertOffsetsToRange(n.start, TextRange.getEnd(n), parseResults.tokenizerOutput.lines),
+        return results.map((r) => ({
+            kind:
+                r.node.nodeType === ParseNodeType.Name && ParseTreeUtils.isWriteAccess(r.node)
+                    ? DocumentHighlightKind.Write
+                    : DocumentHighlightKind.Read,
+            range: convertOffsetsToRange(r.range.start, TextRange.getEnd(r.range), parseResults.tokenizerOutput.lines),
         }));
     }
 }

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

@@ -10,63 +10,123 @@
 
 import { CancellationToken } from 'vscode-languageserver';
 
-import { isCodeUnreachable } from '../analyzer/analyzerNodeInfo';
-import { Declaration } from '../analyzer/declaration';
-import { areDeclarationsSame } from '../analyzer/declarationUtils';
-import { getModuleNode } from '../analyzer/parseTreeUtils';
+import * as AnalyzerNodeInfo from '../analyzer/analyzerNodeInfo';
+import { AliasDeclaration, Declaration, DeclarationType, isAliasDeclaration } from '../analyzer/declaration';
+import {
+    areDeclarationsSame,
+    createSynthesizedAliasDeclaration,
+    getDeclarationsWithUsesLocalNameRemoved,
+} from '../analyzer/declarationUtils';
+import { getModuleNode, getStringNodeValueRange } from '../analyzer/parseTreeUtils';
 import { ParseTreeWalker } from '../analyzer/parseTreeWalker';
+import * as ScopeUtils from '../analyzer/scopeUtils';
+import { isStubFile, SourceMapper } from '../analyzer/sourceMapper';
 import { TypeEvaluator } from '../analyzer/typeEvaluatorTypes';
+import { TypeCategory } from '../analyzer/types';
 import { throwIfCancellationRequested } from '../common/cancellationUtils';
-import { ModuleNameNode, NameNode, ParseNode } from '../parser/parseNodes';
+import { TextRange } from '../common/textRange';
+import { ImportAsNode, NameNode, ParseNode, ParseNodeType, StringNode } from '../parser/parseNodes';
+
+export type CollectionResult = {
+    node: NameNode | StringNode;
+    range: TextRange;
+};
 
 // This walker looks for symbols that are semantically equivalent
 // to the requested symbol.
 export class DocumentSymbolCollector extends ParseTreeWalker {
-    private _symbolName: string;
-    private _declarations: Declaration[] = [];
-    private _startingNode: ParseNode | undefined;
-
-    constructor(
+    static collectFromNode(
         node: NameNode,
-        private _evaluator: TypeEvaluator,
-        private _results: NameNode[],
-        private _cancellationToken: CancellationToken,
-        startingNode?: ParseNode
-    ) {
-        super();
-        this._symbolName = node.value;
+        evaluator: TypeEvaluator,
+        cancellationToken: CancellationToken,
+        startingNode?: ParseNode,
+        treatModuleInImportAndFromImportSame = false
+    ): CollectionResult[] {
+        const symbolName = node.value;
+        const declarations = this.getDeclarationsForNode(
+            node,
+            evaluator,
+            /* resolveLocalName */ true,
+            cancellationToken
+        );
 
-        const declarations = this._evaluator.getDeclarationsForNameNode(node) || [];
+        startingNode = startingNode ?? getModuleNode(node);
+        if (!startingNode) {
+            return [];
+        }
 
+        const collector = new DocumentSymbolCollector(
+            symbolName,
+            declarations,
+            evaluator,
+            cancellationToken,
+            startingNode,
+            treatModuleInImportAndFromImportSame
+        );
+
+        return collector.collect();
+    }
+
+    static getDeclarationsForNode(
+        node: NameNode,
+        evaluator: TypeEvaluator,
+        resolveLocalName: boolean,
+        token: CancellationToken,
+        sourceMapper?: SourceMapper
+    ): Declaration[] {
+        throwIfCancellationRequested(token);
+
+        const declarations = this._getDeclarationsForNode(node, evaluator);
+
+        const resolvedDeclarations: Declaration[] = [];
         declarations.forEach((decl) => {
-            const resolvedDecl = this._evaluator.resolveAliasDeclaration(decl, /* resolveLocalNames */ true);
+            const resolvedDecl = evaluator.resolveAliasDeclaration(decl, resolveLocalName);
             if (resolvedDecl) {
-                this._declarations.push(resolvedDecl);
+                resolvedDeclarations.push(resolvedDecl);
+
+                if (sourceMapper && isStubFile(resolvedDecl.path)) {
+                    const implDecls = sourceMapper.findDeclarations(resolvedDecl);
+                    for (const implDecl of implDecls) {
+                        if (implDecl && implDecl.path) {
+                            this._addIfUnique(resolvedDeclarations, implDecl);
+                        }
+                    }
+                }
             }
         });
 
-        this._startingNode = startingNode ?? getModuleNode(node);
+        return resolvedDeclarations;
+    }
+
+    private _results: CollectionResult[] = [];
+    private _dunderAllNameNodes = new Set<StringNode>();
+
+    constructor(
+        private _symbolName: string,
+        private _declarations: Declaration[],
+        private _evaluator: TypeEvaluator,
+        private _cancellationToken: CancellationToken,
+        private _startingNode: ParseNode,
+        private _treatModuleInImportAndFromImportSame = false
+    ) {
+        super();
+
+        // Don't report strings in __all__ right away, that will
+        // break the assumption on the result ordering.
+        this._setDunderAllNodes(this._startingNode);
     }
 
     collect() {
-        if (!this._startingNode) {
-            return;
-        }
-
         this.walk(this._startingNode);
+        return this._results;
     }
 
     override walk(node: ParseNode) {
-        if (!isCodeUnreachable(node)) {
+        if (!AnalyzerNodeInfo.isCodeUnreachable(node)) {
             super.walk(node);
         }
     }
 
-    override visitModuleName(node: ModuleNameNode): boolean {
-        // Don't ever look for references within a module name.
-        return false;
-    }
-
     override visitName(node: NameNode): boolean {
         throwIfCancellationRequested(this._cancellationToken);
 
@@ -76,7 +136,7 @@ export class DocumentSymbolCollector extends ParseTreeWalker {
         }
 
         if (this._declarations.length > 0) {
-            const declarations = this._evaluator.getDeclarationsForNameNode(node);
+            const declarations = DocumentSymbolCollector._getDeclarationsForNode(node, this._evaluator);
 
             if (declarations && declarations.length > 0) {
                 // Does this name share a declaration with the symbol of interest?
@@ -89,11 +149,22 @@ export class DocumentSymbolCollector extends ParseTreeWalker {
             this._addResult(node);
         }
 
-        return true;
+        return false;
     }
 
-    private _addResult(node: NameNode) {
-        this._results.push(node);
+    override visitString(node: StringNode): boolean {
+        throwIfCancellationRequested(this._cancellationToken);
+
+        if (this._dunderAllNameNodes.has(node)) {
+            this._addResult(node);
+        }
+
+        return false;
+    }
+
+    private _addResult(node: NameNode | StringNode) {
+        const range: TextRange = node.nodeType === ParseNodeType.Name ? node : getStringNodeValueRange(node);
+        this._results.push({ node, range });
     }
 
     private _resultsContainsDeclaration(declaration: Declaration) {
@@ -105,20 +176,211 @@ export class DocumentSymbolCollector extends ParseTreeWalker {
 
         // The reference results declarations are already resolved, so we don't
         // need to call resolveAliasDeclaration on them.
-        if (this._declarations.some((decl) => areDeclarationsSame(decl, resolvedDecl))) {
+        if (
+            this._declarations.some((decl) =>
+                areDeclarationsSame(decl, resolvedDecl, this._treatModuleInImportAndFromImportSame)
+            )
+        ) {
             return true;
         }
 
         // We didn't find the declaration using local-only alias resolution. Attempt
         // it again by fully resolving the alias.
-        const resolvedDeclNonlocal = this._evaluator.resolveAliasDeclaration(
-            resolvedDecl,
-            /* resolveLocalNames */ true
-        );
+        const resolvedDeclNonlocal = this._getResolveAliasDeclaration(resolvedDecl);
         if (!resolvedDeclNonlocal || resolvedDeclNonlocal === resolvedDecl) {
             return false;
         }
 
-        return this._declarations.some((decl) => areDeclarationsSame(decl, resolvedDeclNonlocal));
+        return this._declarations.some((decl) =>
+            areDeclarationsSame(decl, resolvedDeclNonlocal, this._treatModuleInImportAndFromImportSame)
+        );
+    }
+
+    private _getResolveAliasDeclaration(declaration: Declaration) {
+        // TypeEvaluator.resolveAliasDeclaration only resolve alias in AliasDeclaration in the form of
+        // "from x import y as [y]" but don't do thing for alias in "import x as [x]"
+        // Here, alias should have same name as module name.
+        if (isAliasDeclFromImportAsWithAlias(declaration)) {
+            return getDeclarationsWithUsesLocalNameRemoved([declaration])[0];
+        }
+
+        const resolvedDecl = this._evaluator.resolveAliasDeclaration(declaration, /* resolveLocalNames */ true);
+        return isAliasDeclFromImportAsWithAlias(resolvedDecl)
+            ? getDeclarationsWithUsesLocalNameRemoved([resolvedDecl])[0]
+            : resolvedDecl;
+
+        function isAliasDeclFromImportAsWithAlias(decl?: Declaration): decl is AliasDeclaration {
+            return (
+                !!decl &&
+                decl.type === DeclarationType.Alias &&
+                decl.node &&
+                decl.usesLocalName &&
+                decl.node.nodeType === ParseNodeType.ImportAs
+            );
+        }
+    }
+
+    private _setDunderAllNodes(node: ParseNode) {
+        if (node.nodeType !== ParseNodeType.Module) {
+            return;
+        }
+
+        const dunderAllInfo = AnalyzerNodeInfo.getDunderAllInfo(node);
+        if (!dunderAllInfo) {
+            return;
+        }
+
+        const moduleScope = ScopeUtils.getScopeForNode(node);
+        if (!moduleScope) {
+            return;
+        }
+
+        dunderAllInfo.stringNodes.forEach((stringNode) => {
+            if (stringNode.value !== this._symbolName) {
+                return;
+            }
+
+            const symbolInScope = moduleScope.lookUpSymbolRecursive(stringNode.value);
+            if (!symbolInScope) {
+                return;
+            }
+
+            if (!symbolInScope.symbol.getDeclarations().some((d) => this._resultsContainsDeclaration(d))) {
+                return;
+            }
+
+            this._dunderAllNameNodes.add(stringNode);
+        });
+    }
+
+    private static _addIfUnique(declarations: Declaration[], itemToAdd: Declaration) {
+        for (const def of declarations) {
+            if (areDeclarationsSame(def, itemToAdd)) {
+                return;
+            }
+        }
+
+        declarations.push(itemToAdd);
+    }
+
+    private static _getDeclarationsForNode(node: NameNode, evaluator: TypeEvaluator): Declaration[] {
+        // This can handle symbols brought in by wildcard as long as declarations symbol collector
+        // compare against point to actual alias declaration, not one that use local name (ex, import alias)
+        if (node.parent?.nodeType !== ParseNodeType.ModuleName) {
+            let decls = evaluator.getDeclarationsForNameNode(node) || [];
+
+            if (node.parent?.nodeType === ParseNodeType.ImportFromAs) {
+                // Make sure we get the decl for this specific from import statement
+                decls = decls.filter((d) => d.node === node.parent);
+            }
+
+            // If we can't get decl, see whether we can get type from the node.
+            // Some might have synthesized type for the node such as subModule in import X.Y statement.
+            if (decls.length === 0) {
+                const type = evaluator.getType(node);
+                if (type?.category === TypeCategory.Module) {
+                    // Synthesize decl for the module.
+                    return [createSynthesizedAliasDeclaration(type.filePath)];
+                }
+            }
+
+            // We would like to make X in import X and import X.Y as Y to match, but path for
+            // X in import X and one in import X.Y as Y might not match since path in X.Y will point
+            // to X.Y rather than X if import statement has an alias.
+            // so, for such case, we put synthesized one so we can treat X in both statement same.
+            for (const aliasDecl of decls.filter((d) => isAliasDeclaration(d) && !d.loadSymbolsFromPath)) {
+                const node = (aliasDecl as AliasDeclaration).node;
+                if (node.nodeType === ParseNodeType.ImportFromAs) {
+                    // from ... import X case, decl in the submodulefallback has the path.
+                    continue;
+                }
+
+                decls.push(...(evaluator.getDeclarationsForNameNode(node.module.nameParts[0]) || []));
+            }
+
+            return decls;
+        }
+
+        // We treat module name special in find all references. so symbol highlight or rename on multiple files
+        // works even if it is not actually a symbol defined in the file.
+        const moduleName = node.parent;
+        if (
+            moduleName.parent?.nodeType === ParseNodeType.ImportAs ||
+            moduleName.parent?.nodeType === ParseNodeType.ImportFrom
+        ) {
+            const index = moduleName.nameParts.findIndex((n) => n === node);
+
+            // Special case, first module name part.
+            if (index === 0) {
+                // 1. import X or from X import ...
+                let decls: Declaration[] = [];
+
+                // ex, import X as x
+                const isImportAsWithAlias =
+                    moduleName.nameParts.length === 1 &&
+                    moduleName.parent.nodeType === ParseNodeType.ImportAs &&
+                    !!moduleName.parent.alias;
+
+                // if "import" has alias, symbol is assigned to alias, not the module.
+                const importName = isImportAsWithAlias
+                    ? (moduleName.parent as ImportAsNode).alias!.value
+                    : moduleName.nameParts[0].value;
+
+                // First, we need to re-use "decls for X" binder has created
+                // so that it matches with decls type evaluator returns for "references for X".
+                // ex) import X or from .X import ... in init file and etc.
+                const symbolWithScope = ScopeUtils.getScopeForNode(node)?.lookUpSymbolRecursive(importName);
+                if (symbolWithScope) {
+                    decls.push(...symbolWithScope.symbol.getDeclarations().filter((d) => isAliasDeclaration(d)));
+
+                    // If symbols are re-used, then find one that belong to this import statement.
+                    if (decls.length > 1) {
+                        decls = decls.filter((d) => {
+                            d = d as AliasDeclaration;
+
+                            if (d.firstNamePart !== undefined) {
+                                // For multiple import statements with sub modules, decl can be re-used.
+                                // ex) import X.Y and import X.Z or from .X import ... in init file.
+                                // Decls for X will be reused for both import statements, and node will point
+                                // to first import statement. For those case, use firstNamePart instead to check.
+                                return d.firstNamePart === moduleName.nameParts[0].value;
+                            }
+
+                            return d.node === moduleName.parent;
+                        });
+                    }
+
+                    // ex, import X as x
+                    // We have decls for the alias "x" not the module name "X". Convert decls for the "X"
+                    if (isImportAsWithAlias) {
+                        decls = getDeclarationsWithUsesLocalNameRemoved(decls);
+                    }
+                }
+
+                // But, also, we need to put decls for module names type evaluator synthesized so that
+                // we can match both "import X" and "from X import ..."
+                decls.push(
+                    ...(evaluator
+                        .getDeclarationsForNameNode(moduleName.nameParts[0])
+                        ?.filter((d) => isAliasDeclaration(d)) || [])
+                );
+
+                return decls;
+            }
+
+            if (index > 0) {
+                // 2. import X.Y or from X.Y import ....
+                // For submodule "Y", we just use synthesized decls from type evaluator.
+                // Decls for these sub module don't actually exist in the system. Instead, symbol for Y in
+                // "import X.Y" hold onto synthesized module type (without any decl).
+                // And "from X.Y import ..." doesn't have any symbol associated module names.
+                // they can't be referenced in the module.
+                return evaluator.getDeclarationsForNameNode(moduleName.nameParts[index]) || [];
+            }
+
+            return [];
+        }
+
+        return [];
     }
 }

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

@@ -352,7 +352,7 @@ function collectSymbolIndexData(
                 return;
             }
 
-            if (declaration.path.length <= 0) {
+            if (!declaration.loadSymbolsFromPath || declaration.path.length <= 0) {
                 // If alias doesn't have a path to the original file, we can't do dedup
                 // so ignore those aliases.
                 // ex) asyncio.futures, asyncio.base_futures.futures and many will dedup

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

@@ -85,7 +85,11 @@ function _addMissingOptionalToParam(
     const importStatement = importStatements.orderedImports.find((imp) => imp.moduleName === 'typing');
 
     // If there's an existing import statement, insert into it.
-    if (importStatement && importStatement.node.nodeType === ParseNodeType.ImportFrom) {
+    if (
+        importStatement &&
+        importStatement.node.nodeType === ParseNodeType.ImportFrom &&
+        !importStatement.node.isWildcardImport
+    ) {
         const additionalEditActions = getTextEditsForAutoImportSymbolAddition(
             { name: 'Optional' },
             importStatement,

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

@@ -10,19 +10,17 @@
 
 import { CancellationToken } from 'vscode-languageserver';
 
-import * as AnalyzerNodeInfo from '../analyzer/analyzerNodeInfo';
 import { Declaration } from '../analyzer/declaration';
-import * as DeclarationUtils from '../analyzer/declarationUtils';
 import * as ParseTreeUtils from '../analyzer/parseTreeUtils';
-import { ParseTreeWalker } from '../analyzer/parseTreeWalker';
-import { isStubFile, SourceMapper } from '../analyzer/sourceMapper';
+import { SourceMapper } from '../analyzer/sourceMapper';
 import { TypeEvaluator } from '../analyzer/typeEvaluatorTypes';
 import { throwIfCancellationRequested } from '../common/cancellationUtils';
 import { convertOffsetToPosition, convertPositionToOffset } from '../common/positionUtils';
 import { DocumentRange, Position } from '../common/textRange';
 import { TextRange } from '../common/textRange';
-import { ModuleNameNode, NameNode, ParseNode, ParseNodeType } from '../parser/parseNodes';
+import { NameNode, ParseNode, ParseNodeType } from '../parser/parseNodes';
 import { ParseResults } from '../parser/parser';
+import { DocumentSymbolCollector } from './documentSymbolCollector';
 
 export type ReferenceCallback = (locations: DocumentRange[]) => void;
 
@@ -54,9 +52,7 @@ export class ReferencesResult {
     }
 }
 
-export class FindReferencesTreeWalker extends ParseTreeWalker {
-    private readonly _locationsFound: DocumentRange[] = [];
-
+export class FindReferencesTreeWalker {
     constructor(
         private _parseResults: ParseResults,
         private _filePath: string,
@@ -64,91 +60,93 @@ export class FindReferencesTreeWalker extends ParseTreeWalker {
         private _includeDeclaration: boolean,
         private _evaluator: TypeEvaluator,
         private _cancellationToken: CancellationToken
-    ) {
-        super();
-    }
+    ) {}
 
     findReferences(rootNode = this._parseResults.parseTree) {
-        this.walk(rootNode);
+        const collector = new DocumentSymbolCollector(
+            this._referencesResult.symbolName,
+            this._referencesResult.declarations,
+            this._evaluator,
+            this._cancellationToken,
+            rootNode,
+            /* treat module in import and from import same */ true
+        );
 
-        return this._locationsFound;
-    }
-
-    override walk(node: ParseNode) {
-        if (!AnalyzerNodeInfo.isCodeUnreachable(node)) {
-            super.walk(node);
-        }
-    }
-
-    override visitModuleName(node: ModuleNameNode): boolean {
-        // Don't ever look for references within a module name.
-        return false;
-    }
-
-    override visitName(node: NameNode): boolean {
-        throwIfCancellationRequested(this._cancellationToken);
-
-        // No need to do any more work if the symbol name doesn't match.
-        if (node.value !== this._referencesResult.symbolName) {
-            return false;
-        }
-
-        const declarations = this._evaluator.getDeclarationsForNameNode(node);
-
-        if (declarations && declarations.length > 0) {
-            // Does this name share a declaration with the symbol of interest?
-            if (declarations.some((decl) => this._resultsContainsDeclaration(decl))) {
-                // Is it the same symbol?
-                if (this._includeDeclaration || node !== this._referencesResult.nodeAtOffset) {
-                    this._locationsFound.push({
-                        path: this._filePath,
-                        range: {
-                            start: convertOffsetToPosition(node.start, this._parseResults.tokenizerOutput.lines),
-                            end: convertOffsetToPosition(
-                                TextRange.getEnd(node),
-                                this._parseResults.tokenizerOutput.lines
-                            ),
-                        },
-                    });
-                }
+        const results: DocumentRange[] = [];
+        for (const result of collector.collect()) {
+            // Is it the same symbol?
+            if (this._includeDeclaration || result.node !== this._referencesResult.nodeAtOffset) {
+                results.push({
+                    path: this._filePath,
+                    range: {
+                        start: convertOffsetToPosition(result.range.start, this._parseResults.tokenizerOutput.lines),
+                        end: convertOffsetToPosition(
+                            TextRange.getEnd(result.range),
+                            this._parseResults.tokenizerOutput.lines
+                        ),
+                    },
+                });
             }
         }
 
-        return true;
-    }
-
-    private _resultsContainsDeclaration(declaration: Declaration) {
-        // Resolve the declaration.
-        const resolvedDecl = this._evaluator.resolveAliasDeclaration(declaration, /* resolveLocalNames */ false);
-        if (!resolvedDecl) {
-            return false;
-        }
-
-        // The reference results declarations are already resolved, so we don't
-        // need to call resolveAliasDeclaration on them.
-        if (
-            this._referencesResult.declarations.some((decl) => DeclarationUtils.areDeclarationsSame(decl, resolvedDecl))
-        ) {
-            return true;
-        }
-
-        // We didn't find the declaration using local-only alias resolution. Attempt
-        // it again by fully resolving the alias.
-        const resolvedDeclNonlocal = this._evaluator.resolveAliasDeclaration(
-            resolvedDecl,
-            /* resolveLocalNames */ true
-        );
-        if (!resolvedDeclNonlocal || resolvedDeclNonlocal === resolvedDecl) {
-            return false;
-        }
-
-        return this._referencesResult.declarations.some((decl) =>
-            DeclarationUtils.areDeclarationsSame(decl, resolvedDeclNonlocal)
-        );
+        return results;
     }
 }
 
 export class ReferencesProvider {
+    static getDeclarationForNode(
+        sourceMapper: SourceMapper,
+        filePath: string,
+        node: NameNode,
+        evaluator: TypeEvaluator,
+        reporter: ReferenceCallback | undefined,
+        token: CancellationToken
+    ) {
+        throwIfCancellationRequested(token);
+
+        const declarations = DocumentSymbolCollector.getDeclarationsForNode(
+            node,
+            evaluator,
+            /* resolveLocalNames */ false,
+            token,
+            sourceMapper
+        );
+
+        if (declarations.length === 0) {
+            return undefined;
+        }
+
+        // Does this symbol require search beyond the current file? Determine whether
+        // the symbol is declared within an evaluation scope that is within the current
+        // file and cannot be imported directly from other modules.
+        const requiresGlobalSearch = declarations.some((decl) => {
+            // If the declaration is outside of this file, a global search is needed.
+            if (decl.path !== filePath) {
+                return true;
+            }
+
+            const evalScope = ParseTreeUtils.getEvaluationScopeNode(decl.node);
+
+            // If the declaration is at the module level or a class level, it can be seen
+            // outside of the current module, so a global search is needed.
+            if (evalScope.nodeType === ParseNodeType.Module || evalScope.nodeType === ParseNodeType.Class) {
+                return true;
+            }
+
+            // If the name node is a member variable, we need to do a global search.
+            if (
+                decl.node?.parent?.nodeType === ParseNodeType.MemberAccess &&
+                decl.node === decl.node.parent.memberName
+            ) {
+                return true;
+            }
+
+            return false;
+        });
+
+        return new ReferencesResult(requiresGlobalSearch, node, node.value, declarations, reporter);
+    }
+
     static getDeclarationForPosition(
         sourceMapper: SourceMapper,
         parseResults: ParseResults,
@@ -175,76 +173,7 @@ export class ReferencesProvider {
             return undefined;
         }
 
-        // Special case module names, which don't have references.
-        if (node.parent?.nodeType === ParseNodeType.ModuleName) {
-            return undefined;
-        }
-
-        const declarations = evaluator.getDeclarationsForNameNode(node);
-        if (!declarations) {
-            return undefined;
-        }
-
-        const resolvedDeclarations: Declaration[] = [];
-        declarations.forEach((decl) => {
-            const resolvedDecl = evaluator.resolveAliasDeclaration(decl, /* resolveLocalNames */ false);
-            if (resolvedDecl) {
-                resolvedDeclarations.push(resolvedDecl);
-
-                if (isStubFile(resolvedDecl.path)) {
-                    const implDecls = sourceMapper.findDeclarations(resolvedDecl);
-                    for (const implDecl of implDecls) {
-                        if (implDecl && implDecl.path) {
-                            this._addIfUnique(resolvedDeclarations, implDecl);
-                        }
-                    }
-                }
-            }
-        });
-
-        if (resolvedDeclarations.length === 0) {
-            return undefined;
-        }
-
-        // Does this symbol require search beyond the current file? Determine whether
-        // the symbol is declared within an evaluation scope that is within the current
-        // file and cannot be imported directly from other modules.
-        const requiresGlobalSearch = resolvedDeclarations.some((decl) => {
-            // If the declaration is outside of this file, a global search is needed.
-            if (decl.path !== filePath) {
-                return true;
-            }
-
-            const evalScope = ParseTreeUtils.getEvaluationScopeNode(decl.node);
-
-            // If the declaration is at the module level or a class level, it can be seen
-            // outside of the current module, so a global search is needed.
-            if (evalScope.nodeType === ParseNodeType.Module || evalScope.nodeType === ParseNodeType.Class) {
-                return true;
-            }
-
-            // If the name node is a member variable, we need to do a global search.
-            if (
-                decl.node?.parent?.nodeType === ParseNodeType.MemberAccess &&
-                decl.node === decl.node.parent.memberName
-            ) {
-                return true;
-            }
-
-            return false;
-        });
-
-        return new ReferencesResult(requiresGlobalSearch, node, node.value, resolvedDeclarations, reporter);
-    }
-
-    private static _addIfUnique(declarations: Declaration[], itemToAdd: Declaration) {
-        for (const def of declarations) {
-            if (DeclarationUtils.areDeclarationsSame(def, itemToAdd)) {
-                return;
-            }
-        }
-
-        declarations.push(itemToAdd);
+        return this.getDeclarationForNode(sourceMapper, filePath, node, evaluator, reporter, token);
     }
 
     static addReferences(

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

@@ -0,0 +1,671 @@
+/*
+ * renameModuleProvider.ts
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ *
+ * Logic that updates affected references of a module rename/move.
+ */
+
+import { CancellationToken } from 'vscode-languageserver';
+
+import { AliasDeclaration, isAliasDeclaration } from '../analyzer/declaration';
+import { createSynthesizedAliasDeclaration } from '../analyzer/declarationUtils';
+import { createImportedModuleDescriptor, ImportResolver, ModuleNameAndType } from '../analyzer/importResolver';
+import {
+    getImportGroupFromModuleNameAndType,
+    getTextEditsForAutoImportInsertion,
+    getTextEditsForAutoImportSymbolAddition,
+    getTextRangeForImportNameDeletion,
+    getTopLevelImports,
+    ImportNameInfo,
+    ImportStatements,
+} from '../analyzer/importStatementUtils';
+import {
+    getDottedNameWithGivenNodeAsLastName,
+    getFirstAncestorOrSelfOfKind,
+    getFullStatementRange,
+    isFromImportAlias,
+    isFromImportModuleName,
+    isFromImportName,
+    isImportAlias,
+    isImportModuleName,
+    isLastNameOfModuleName,
+} from '../analyzer/parseTreeUtils';
+import { isStubFile } from '../analyzer/sourceMapper';
+import { TypeEvaluator } from '../analyzer/typeEvaluatorTypes';
+import { getOrAdd, removeArrayElements } from '../common/collectionUtils';
+import { ConfigOptions } from '../common/configOptions';
+import { isString } from '../common/core';
+import { FileEditAction } from '../common/editAction';
+import { FileSystem } from '../common/fileSystem';
+import {
+    getDirectoryPath,
+    getFileName,
+    getRelativePathComponentsFromDirectory,
+    stripFileExtension,
+} from '../common/pathUtils';
+import { convertOffsetToPosition, convertTextRangeToRange } from '../common/positionUtils';
+import { doRangesIntersect, extendRange, Range, rangesAreEqual, TextRange } from '../common/textRange';
+import {
+    ImportAsNode,
+    ImportFromAsNode,
+    ImportFromNode,
+    ModuleNameNode,
+    NameNode,
+    ParseNodeType,
+} from '../parser/parseNodes';
+import { ParseResults } from '../parser/parser';
+import { DocumentSymbolCollector } from './documentSymbolCollector';
+
+export class RenameModuleProvider {
+    static create(
+        importResolver: ImportResolver,
+        configOptions: ConfigOptions,
+        evaluator: TypeEvaluator,
+        moduleFilePath: string,
+        newModuleFilePath: string,
+        token: CancellationToken
+    ) {
+        const execEnv = configOptions.findExecEnvironment(moduleFilePath);
+        const moduleName = importResolver.getModuleNameForImport(moduleFilePath, execEnv);
+        if (!moduleName.moduleName) {
+            return undefined;
+        }
+
+        const newModuleName = importResolver.getModuleNameForImport(newModuleFilePath, execEnv);
+        if (!newModuleName.moduleName) {
+            return undefined;
+        }
+
+        // Create synthesized alias decls from the given file path. If the given file is for stub,
+        // create one for the corresponding py file as well.
+        const moduleDecls = [createSynthesizedAliasDeclaration(moduleFilePath)];
+        if (isStubFile(moduleFilePath)) {
+            // The resolveImport should make sure non stub file search to happen.
+            importResolver.resolveImport(
+                moduleFilePath,
+                execEnv,
+                createImportedModuleDescriptor(moduleName.moduleName)
+            );
+
+            importResolver
+                .getSourceFilesFromStub(moduleFilePath, execEnv, /*mapCompiled*/ false)
+                .forEach((p) => moduleDecls.push(createSynthesizedAliasDeclaration(p)));
+        }
+
+        return new RenameModuleProvider(
+            importResolver.fileSystem,
+            evaluator,
+            newModuleFilePath,
+            moduleName,
+            newModuleName,
+            moduleDecls,
+            token
+        );
+    }
+
+    private readonly _moduleNames: string[];
+    private readonly _newModuleNames: string[];
+    private readonly _onlyNameChanged: boolean;
+    private readonly _results = new Map<string, FileEditAction[]>();
+
+    private readonly _aliasIntroduced = new Set<ImportAsNode>();
+
+    private constructor(
+        private _fs: FileSystem,
+        private _evaluator: TypeEvaluator,
+        private _newModuleFilePath: string,
+        private _moduleNameAndType: ModuleNameAndType,
+        private _newModuleNameAndType: ModuleNameAndType,
+        private _moduleDecls: AliasDeclaration[],
+        private _token: CancellationToken
+    ) {
+        // moduleName and newModuleName are always in the absolute path form.
+        this._moduleNames = this._moduleName.split('.');
+        this._newModuleNames = this._newModuleName.split('.');
+
+        if (this._moduleNames.length !== this._newModuleNames.length) {
+            this._onlyNameChanged = false;
+            return;
+        }
+
+        let i = 0;
+        for (i = 0; i < this._moduleNames.length - 1; i++) {
+            if (this._moduleNames[i] !== this._newModuleNames[i]) {
+                break;
+            }
+        }
+
+        this._onlyNameChanged = i === this._moduleNames.length - 1;
+    }
+
+    renameModuleReferences(filePath: string, parseResults: ParseResults) {
+        const collector = new DocumentSymbolCollector(
+            this.symbolName,
+            this._moduleDecls,
+            this._evaluator!,
+            this._token,
+            parseResults.parseTree,
+            /*treatModuleImportAndFromImportSame*/ true
+        );
+
+        const results = collector.collect();
+
+        const nameRemoved = new Set<NameNode>();
+        let importStatements: ImportStatements | undefined;
+        for (const result of results) {
+            const nodeFound = result.node;
+
+            if (nodeFound.nodeType === ParseNodeType.String) {
+                // ex) __all__ = ["[a]"]
+                this._addResultWithTextRange(filePath, result.range, parseResults, this._newSymbolName);
+                continue;
+            }
+
+            if (isImportModuleName(nodeFound)) {
+                if (!isLastNameOfModuleName(nodeFound)) {
+                    // It must be directory and we don't support folder rename.
+                    continue;
+                }
+
+                const moduleNameNode = getFirstAncestorOrSelfOfKind(
+                    nodeFound,
+                    ParseNodeType.ModuleName
+                ) as ModuleNameNode;
+
+                // * Enhancement * one case we don't handle is introducing new symbol in __all__
+                // or converting "import" statement to "from import" statement.
+                //
+                // when the existing statement was "import x as x" and it is changed to
+                // "import y.z as z". we either need to introduce "z" in __all__ or convert
+                // "import y.z as z" to "from y import z as z" to make sure we keep the symbol
+                // visibility same.
+                //
+                // when we convert "import x as x" to "from y import z as z", we need to handle
+                // deletion of existing import statement or (x as x) and inserting/merging
+                // new "from import" statement.
+
+                // If original module name was single word and it becomes dotted name without alias,
+                // then we introduce alias to keep references as a single word.
+                // ex) import [xxx] to import [aaa.bbb as bbb]
+                if (
+                    moduleNameNode.nameParts.length === 1 &&
+                    moduleNameNode.parent?.nodeType === ParseNodeType.ImportAs &&
+                    !moduleNameNode.parent.alias &&
+                    this._newModuleNames.length > 1
+                ) {
+                    this._aliasIntroduced.add(moduleNameNode.parent);
+
+                    this._addResultWithTextRange(
+                        filePath,
+                        moduleNameNode,
+                        parseResults,
+                        `${this._newModuleName} as ${this._newSymbolName}`
+                    );
+                    continue;
+                }
+
+                // Otherwise, update whole module name to new name
+                // ex) import [xxx.yyy] to import [aaa.bbb]
+                this._addResultWithTextRange(filePath, moduleNameNode, parseResults, this._newModuleName);
+                continue;
+            }
+
+            if (isImportAlias(nodeFound)) {
+                // ex) import xxx as [yyy] to import xxx as [zzz]
+                this._addResultWithTextRange(filePath, result.range, parseResults, this._newSymbolName);
+                continue;
+            }
+
+            if (isFromImportModuleName(nodeFound)) {
+                if (!isLastNameOfModuleName(nodeFound)) {
+                    // It must be directory and we don't support folder rename.
+                    continue;
+                }
+
+                const moduleNameNode = getFirstAncestorOrSelfOfKind(
+                    nodeFound,
+                    ParseNodeType.ModuleName
+                ) as ModuleNameNode;
+
+                const fromNode = moduleNameNode.parent as ImportFromNode;
+
+                // We need to check whether imports of this import statement has
+                // any implicit submodule imports or not. If there is one, we need to
+                // either split or leave it as it is.
+                const exportedSymbols = [];
+                const subModules = [];
+                for (const importFromAs of fromNode.imports) {
+                    if (this._isExportedSymbol(importFromAs.name)) {
+                        exportedSymbols.push(importFromAs);
+                    } else {
+                        subModules.push(importFromAs);
+                    }
+                }
+
+                if (subModules.length === 0) {
+                    // We don't have any sub modules, we can change module name to new one.
+                    // Update whole module name to new name.
+                    // ex) from [xxx.yyy] import zzz to from [aaa.bbb] import zzz
+                    this._addResultWithTextRange(filePath, moduleNameNode, parseResults, this._newModuleName);
+                    continue;
+                }
+
+                if (exportedSymbols.length === 0) {
+                    // We only have sub modules. That means module name actually refers to
+                    // folder name, not module (ex, __init__.py). Since we don't support
+                    // renaming folder, leave things as they are.
+                    continue;
+                }
+
+                // Now, we need to split "from import" statement to 2.
+
+                // First, delete existing exported symbols from "from import" statement.
+                for (const importFromAs of exportedSymbols) {
+                    this._addImportNameDeletion(filePath, parseResults, nameRemoved, fromNode.imports, importFromAs);
+                }
+
+                importStatements =
+                    importStatements ?? getTopLevelImports(parseResults.parseTree, /*includeImplicitImports*/ false);
+
+                // For now, this won't merge absolute and relative path "from import"
+                // statement.
+                this._addResultEdits(
+                    this._getTextEditsForNewOrExistingFromImport(
+                        filePath,
+                        fromNode,
+                        parseResults,
+                        nameRemoved,
+                        importStatements,
+                        this._newModuleName,
+                        exportedSymbols.map((i) => {
+                            const name =
+                                results.findIndex((r) => r.node === i.name) >= 0 ? this._newSymbolName : i.name.value;
+                            const alias =
+                                results.findIndex((r) => r.node === i.alias) >= 0
+                                    ? this._newSymbolName
+                                    : i.alias?.value;
+
+                            return { name, alias };
+                        })
+                    )
+                );
+                continue;
+            }
+
+            if (isFromImportName(nodeFound)) {
+                if (nameRemoved.has(nodeFound)) {
+                    // Import name is already removed.
+                    continue;
+                }
+
+                const fromNode = nodeFound.parent?.parent as ImportFromNode;
+                const newModuleName = this._getNewModuleName(filePath, fromNode.module);
+
+                // If the name bound to symbol re-exported, we don't need to update module name.
+                // Existing logic should make sure re-exported symbol name work as before after
+                // symbol rename.
+                if (this._isExportedSymbol(nodeFound)) {
+                    this._addResultWithTextRange(filePath, result.range, parseResults, this._newSymbolName);
+                    continue;
+                }
+
+                if (fromNode.imports.length === 1) {
+                    // ex) from xxx import [yyy] to from [aaa.bbb] import [zzz]
+                    this._addResultWithTextRange(filePath, fromNode.module, parseResults, newModuleName);
+                    this._addResultWithTextRange(filePath, result.range, parseResults, this._newSymbolName);
+                } else {
+                    // Delete the existing import name including alias.
+                    const importFromAs = nodeFound.parent as ImportFromAsNode;
+                    this._addImportNameDeletion(filePath, parseResults, nameRemoved, fromNode.imports, importFromAs);
+
+                    importStatements =
+                        importStatements ??
+                        getTopLevelImports(parseResults.parseTree, /*includeImplicitImports*/ false);
+
+                    // ex) from xxx import yyy, [zzz] to
+                    //     from xxx import yyy
+                    //     from [aaa.bbb] import [ccc]
+                    // or
+                    //     from aaa.bbb import ddd
+                    //     from xxx import yyy, [zzz] to
+                    //     from aaa.bbb import [ccc], ddd
+                    //
+                    // For now, this won't merge absolute and relative path "from import"
+                    // statement.
+                    const importNameInfo = {
+                        name: this._newSymbolName,
+                        alias:
+                            importFromAs.alias?.value === this.symbolName
+                                ? this._newSymbolName
+                                : importFromAs.alias?.value,
+                    };
+
+                    this._addResultEdits(
+                        this._getTextEditsForNewOrExistingFromImport(
+                            filePath,
+                            fromNode,
+                            parseResults,
+                            nameRemoved,
+                            importStatements,
+                            newModuleName,
+                            [importNameInfo]
+                        )
+                    );
+                }
+                continue;
+            }
+
+            if (isFromImportAlias(nodeFound)) {
+                if (nameRemoved.has(nodeFound)) {
+                    // alias is already removed.
+                    continue;
+                }
+
+                // ex) from ccc import xxx as [yyy] to from ccc import xxx as [zzz]
+                this._addResultWithTextRange(filePath, result.range, parseResults, this._newSymbolName);
+                continue;
+            }
+
+            /** TODO: if we get more than 1 decls, flag it as attention needed */
+            const decls = DocumentSymbolCollector.getDeclarationsForNode(
+                nodeFound,
+                this._evaluator,
+                /*resolveLocalName*/ false,
+                this._token
+            ).filter((d) => isAliasDeclaration(d)) as AliasDeclaration[];
+
+            if (this._onlyNameChanged) {
+                // Simple case. only name has changed. but not path.
+                // Just replace name to new symbol name.
+                // ex) a.[b].foo() to a.[z].foo()
+                this._addResultWithTextRange(filePath, result.range, parseResults, this._newSymbolName);
+                continue;
+            }
+
+            if (
+                decls?.some(
+                    (d) =>
+                        !d.usesLocalName &&
+                        (!d.node || d.node.nodeType === ParseNodeType.ImportAs) &&
+                        !this._aliasIntroduced.has(d.node)
+                )
+            ) {
+                const dottedName = getDottedNameWithGivenNodeAsLastName(nodeFound);
+                if (dottedName.parent?.nodeType !== ParseNodeType.MemberAccess) {
+                    // Replace whole dotted name with new module name.
+                    this._addResultWithTextRange(filePath, dottedName, parseResults, this._newModuleName);
+                    continue;
+                }
+
+                // Check whether name after me is sub module or not.
+                // ex) a.b.[c]
+                const nextNameDecl = this._evaluator.getDeclarationsForNameNode(dottedName.parent.memberName);
+                if (!nextNameDecl || nextNameDecl.length === 0) {
+                    // Next dotted name is sub module. That means dottedName actually refers to folder names, not modules.
+                    // and We don't support renaming folder. So, leave things as they are.
+                    // ex) import a.b.c
+                    //     [a.b].[c]
+                    continue;
+                }
+
+                // Next name is actual symbol. Replace whole name to new module name.
+                // ex) import a.b.c
+                //     [a.b.c].[foo]()
+                this._addResultWithTextRange(filePath, dottedName, parseResults, this._newModuleName);
+                continue;
+            }
+
+            if (result.node.value !== this._newSymbolName) {
+                this._addResultWithTextRange(filePath, result.range, parseResults, this._newSymbolName);
+                continue;
+            }
+        }
+    }
+
+    private _isExportedSymbol(nameNode: NameNode): boolean {
+        const decls = this._evaluator.getDeclarationsForNameNode(nameNode);
+        if (!decls) {
+            return false;
+        }
+
+        // If submoduleFallback exists, then, it points to submodule not symbol.
+        return !decls.some((d) => isAliasDeclaration(d) && d.submoduleFallback);
+    }
+
+    private _getNewModuleName(currentFilePath: string, moduleName: ModuleNameNode) {
+        if (moduleName.leadingDots === 0) {
+            const newModuleName = this._newModuleName.substr(
+                0,
+                this._newModuleName.length - this._newSymbolName.length - 1
+            );
+            return newModuleName.length > 0 ? newModuleName : '.';
+        }
+
+        // If the existing code was using relative path, try to keep the relative path.
+        const relativePaths = getRelativePathComponentsFromDirectory(
+            getDirectoryPath(currentFilePath),
+            getDirectoryPath(this._newModuleFilePath),
+            (f) => this._fs.realCasePath(f)
+        );
+
+        // Both file paths are pointing to user files. So we don't need to worry about
+        // relative path pointing to library files.
+        let relativeModuleName = '.';
+        for (let i = 1; i < relativePaths.length; i++) {
+            const relativePath = relativePaths[i];
+            if (relativePath === '..') {
+                relativeModuleName += '.';
+            } else {
+                relativeModuleName += relativePath;
+            }
+
+            if (relativePath !== '..' && i !== relativePaths.length - 1) {
+                relativeModuleName += '.';
+            }
+        }
+
+        // __init__ makes the folder itself not file inside of the folder part of
+        // module path. Move up one more level.
+        const fileName = stripFileExtension(getFileName(this._newModuleFilePath));
+        if (fileName === '__init__') {
+            relativeModuleName += '.';
+        }
+
+        return relativeModuleName;
+    }
+
+    getEdits(): FileEditAction[] {
+        const edits: FileEditAction[] = [];
+        this._results.forEach((v) => edits.push(...v));
+
+        return edits;
+    }
+
+    get symbolName() {
+        return this._moduleNames[this._moduleNames.length - 1];
+    }
+
+    private get _moduleName() {
+        return this._moduleNameAndType.moduleName;
+    }
+
+    private get _newModuleName() {
+        return this._newModuleNameAndType.moduleName;
+    }
+
+    private get _newSymbolName() {
+        return this._newModuleNames[this._newModuleNames.length - 1];
+    }
+
+    private _addImportNameDeletion(
+        filePath: string,
+        parseResults: ParseResults,
+        nameRemoved: Set<NameNode>,
+        imports: ImportFromAsNode[],
+        importToDelete: ImportFromAsNode
+    ) {
+        const range = getTextRangeForImportNameDeletion(
+            imports,
+            imports.findIndex((v) => v === importToDelete)
+        );
+
+        this._addResultWithTextRange(filePath, range, parseResults, '');
+
+        // Mark that we don't need to process these node again later.
+        nameRemoved.add(importToDelete.name);
+        if (importToDelete.alias) {
+            nameRemoved.add(importToDelete.alias);
+        }
+
+        // Check whether we have deleted all trailing import names.
+        // If either no trailing import is deleted or handled properly
+        // then, there is nothing to do. otherwise, either delete the whole statement
+        // or remove trailing comma.
+        // ex) from x import [y], z or from x import y[, z]
+        let lastImportIndexNotDeleted = 0;
+        for (
+            lastImportIndexNotDeleted = imports.length - 1;
+            lastImportIndexNotDeleted >= 0;
+            lastImportIndexNotDeleted--
+        ) {
+            if (!nameRemoved.has(imports[lastImportIndexNotDeleted].name)) {
+                break;
+            }
+        }
+
+        if (lastImportIndexNotDeleted === -1) {
+            // Whole statement is deleted. Remove the statement itself.
+            // ex) [from x import a, b, c]
+            const fromImport = getFirstAncestorOrSelfOfKind(importToDelete, ParseNodeType.ImportFrom);
+            if (fromImport) {
+                this._addResultWithRange(filePath, getFullStatementRange(fromImport, parseResults.tokenizerOutput), '');
+            }
+        } else if (lastImportIndexNotDeleted >= 0 && lastImportIndexNotDeleted < imports.length - 2) {
+            // We need to delete trailing comma
+            // ex) from x import a, [b, c]
+            const start = TextRange.getEnd(imports[lastImportIndexNotDeleted]);
+            const length = TextRange.getEnd(imports[lastImportIndexNotDeleted + 1]) - start;
+            this._addResultWithTextRange(filePath, { start, length }, parseResults, '');
+        }
+    }
+
+    private _addResultWithTextRange(filePath: string, range: TextRange, parseResults: ParseResults, newName: string) {
+        const existing = parseResults.text.substr(range.start, range.length);
+        if (existing === newName) {
+            // No change. Return as it is.
+            return;
+        }
+
+        this._addResultWithRange(filePath, convertTextRangeToRange(range, parseResults.tokenizerOutput.lines), newName);
+    }
+
+    private _addResultEdits(edits: FileEditAction[]) {
+        edits.forEach((e) => this._addResultWithRange(e.filePath, e.range, e.replacementText));
+    }
+
+    private _getDeletionsForSpan(filePathOrEdit: string | FileEditAction[], range: Range) {
+        if (isString(filePathOrEdit)) {
+            filePathOrEdit = this._results.get(filePathOrEdit) ?? [];
+        }
+
+        return filePathOrEdit.filter((e) => e.replacementText === '' && doRangesIntersect(e.range, range));
+    }
+
+    private _removeEdits(filePathOrEdit: string | FileEditAction[], edits: FileEditAction[]) {
+        if (isString(filePathOrEdit)) {
+            filePathOrEdit = this._results.get(filePathOrEdit) ?? [];
+        }
+
+        removeArrayElements(filePathOrEdit, (f) => edits.findIndex((e) => e === f) >= 0);
+    }
+
+    private _addResultWithRange(filePath: string, range: Range, replacementText: string) {
+        const edits = getOrAdd(this._results, filePath, () => []);
+        if (replacementText === '') {
+            // If it is a deletion, merge with overlapping deletion edit if there is any.
+            const deletions = this._getDeletionsForSpan(edits, range);
+            if (deletions.length > 0) {
+                // Delete the existing ones.
+                this._removeEdits(edits, deletions);
+
+                // Extend range with deleted ones.
+                extendRange(
+                    range,
+                    deletions.map((d) => d.range)
+                );
+            }
+        }
+
+        // Don't put duplicated edit. It can happen if code has duplicated module import.
+        // ex) from a import b, b, c
+        // If we need to introduce new "from import" statement for "b", we will add new statement twice.
+        if (edits.some((e) => rangesAreEqual(e.range, range) && e.replacementText === replacementText)) {
+            return;
+        }
+
+        edits.push({ filePath, range, replacementText });
+    }
+
+    private _getTextEditsForNewOrExistingFromImport(
+        filePath: string,
+        currentFromImport: ImportFromNode,
+        parseResults: ParseResults,
+        nameRemoved: Set<NameNode>,
+        importStatements: ImportStatements,
+        moduleName: string,
+        importNameInfo: ImportNameInfo[]
+    ): FileEditAction[] {
+        // See whether we have existing from import statement for the same module
+        // ex) from [|moduleName|] import subModule
+        const imported = importStatements.orderedImports.find((i) => i.moduleName === moduleName);
+        if (imported && imported.node.nodeType === ParseNodeType.ImportFrom && !imported.node.isWildcardImport) {
+            const edits = getTextEditsForAutoImportSymbolAddition(importNameInfo, imported, parseResults);
+            if (imported.node !== currentFromImport) {
+                // Add what we want to the existing "import from" statement as long as it is not the same import
+                // node we are working on.
+                return edits.map((e) => ({ filePath, range: e.range, replacementText: e.replacementText }));
+            }
+
+            // Check whether we can avoid creating a new statement. We can't just merge with existing one since
+            // we could create invalid text edits (2 edits that change the same span, or invalid replacement text since
+            // texts on the node has changed)
+            if (this._onlyNameChanged && importNameInfo.length === 1 && edits.length === 1) {
+                const deletions = this._getDeletionsForSpan(filePath, edits[0].range);
+                if (deletions.length === 0) {
+                    return [{ filePath, range: edits[0].range, replacementText: edits[0].replacementText }];
+                } else {
+                    const alias =
+                        importNameInfo[0].alias === this._newSymbolName ? this.symbolName : importNameInfo[0].alias;
+
+                    const importName = currentFromImport.imports.find(
+                        (i) => i.name.value === this.symbolName && i.alias?.value === alias
+                    );
+                    if (importName) {
+                        this._removeEdits(filePath, deletions);
+                        if (importName.alias) {
+                            nameRemoved.delete(importName.alias);
+                        }
+
+                        return [
+                            {
+                                filePath,
+                                range: convertTextRangeToRange(importName.name, parseResults.tokenizerOutput.lines),
+                                replacementText: this._newSymbolName,
+                            },
+                        ];
+                    }
+                }
+            }
+        }
+
+        return getTextEditsForAutoImportInsertion(
+            importNameInfo,
+            importStatements,
+            moduleName,
+            getImportGroupFromModuleNameAndType(this._newModuleNameAndType),
+            parseResults,
+            convertOffsetToPosition(parseResults.parseTree.length, parseResults.tokenizerOutput.lines)
+        ).map((e) => ({ filePath, range: e.range, replacementText: e.replacementText }));
+    }
+}

packages/pyright-internal/src/pyrightFileSystem.ts

@@ -13,36 +13,18 @@ import type * as fs from 'fs';
 
 import { getPyTypedInfo } from './analyzer/pyTypedUtils';
 import { ExecutionEnvironment } from './common/configOptions';
-import {
-    FileSystem,
-    FileWatcher,
-    FileWatcherEventHandler,
-    MkDirOptions,
-    Stats,
-    TmpfileOptions,
-    VirtualDirent,
-} from './common/fileSystem';
+import { FileSystem, MkDirOptions } from './common/fileSystem';
 import { stubsSuffix } from './common/pathConsts';
 import {
     changeAnyExtension,
     combinePaths,
     ensureTrailingDirectorySeparator,
-    getDirectoryPath,
-    getFileName,
     isDirectory,
     tryStat,
 } from './common/pathUtils';
+import { ReadOnlyAugmentedFileSystem } from './readonlyAugmentedFileSystem';
 
-export class PyrightFileSystem implements FileSystem {
-    // Mapped file to original file map
-    private readonly _fileMap = new Map<string, string>();
-
-    // Original file to mapped file map
-    private readonly _reverseFileMap = new Map<string, string>();
-
-    // Mapped files per a containing folder map
-    private readonly _folderMap = new Map<string, string[]>();
-
+export class PyrightFileSystem extends ReadOnlyAugmentedFileSystem {
     // Root paths processed
     private readonly _rootSearched = new Set<string>();
 
@@ -55,118 +37,35 @@ export class PyrightFileSystem implements FileSystem {
 
     private readonly _customUriMap = new Map<string, { uri: string; closed: boolean; hasPendingRequest: boolean }>();
 
-    constructor(private _realFS: FileSystem) {}
-
-    existsSync(path: string): boolean {
-        if (this._isVirtualEntry(path)) {
-            // Pretend partial stub folder and its files not exist
-            return false;
-        }
-
-        return this._realFS.existsSync(this._getPartialStubOriginalPath(path));
+    constructor(realFS: FileSystem) {
+        super(realFS);
     }
 
-    mkdirSync(path: string, options?: MkDirOptions): void {
+    override mkdirSync(path: string, options?: MkDirOptions): void {
         this._realFS.mkdirSync(path, options);
     }
 
-    chdir(path: string): void {
+    override chdir(path: string): void {
         this._realFS.chdir(path);
     }
 
-    readdirEntriesSync(path: string): fs.Dirent[] {
-        const entries = this._realFS.readdirEntriesSync(path).filter((item) => {
-            // Filter out the stub package directory.
-            return !this._isVirtualEntry(combinePaths(path, item.name));
-        });
-
-        const partialStubs = this._folderMap.get(ensureTrailingDirectorySeparator(path));
-        if (!partialStubs) {
-            return entries;
-        }
-
-        return entries.concat(partialStubs.map((f) => new VirtualDirent(f, /* file */ true)));
+    override writeFileSync(path: string, data: string | Buffer, encoding: BufferEncoding | null): void {
+        this._realFS.writeFileSync(this._getOriginalPath(path), data, encoding);
     }
 
-    readdirSync(path: string): string[] {
-        const entries = this._realFS.readdirSync(path).filter((item) => {
-            // Filter out the stub package directory.
-            return !this._isVirtualEntry(combinePaths(path, item));
-        });
-
-        const partialStubs = this._folderMap.get(ensureTrailingDirectorySeparator(path));
-        if (!partialStubs) {
-            return entries;
-        }
-
-        return entries.concat(partialStubs);
+    override unlinkSync(path: string): void {
+        this._realFS.unlinkSync(this._getOriginalPath(path));
     }
 
-    readFileSync(path: string, encoding?: null): Buffer;
-    readFileSync(path: string, encoding: BufferEncoding): string;
-    readFileSync(path: string, encoding?: BufferEncoding | null): string | Buffer {
-        return this._realFS.readFileSync(this._getPartialStubOriginalPath(path), encoding);
+    override createWriteStream(path: string): fs.WriteStream {
+        return this._realFS.createWriteStream(this._getOriginalPath(path));
     }
 
-    writeFileSync(path: string, data: string | Buffer, encoding: BufferEncoding | null): void {
-        this._realFS.writeFileSync(this._getPartialStubOriginalPath(path), data, encoding);
+    override copyFileSync(src: string, dst: string): void {
+        this._realFS.copyFileSync(this._getOriginalPath(src), this._getOriginalPath(dst));
     }
 
-    statSync(path: string): Stats {
-        return this._realFS.statSync(this._getPartialStubOriginalPath(path));
-    }
-
-    unlinkSync(path: string): void {
-        this._realFS.unlinkSync(this._getPartialStubOriginalPath(path));
-    }
-
-    realpathSync(path: string): string {
-        return this._realFS.realpathSync(path);
-    }
-
-    getModulePath(): string {
-        return this._realFS.getModulePath();
-    }
-
-    createFileSystemWatcher(paths: string[], listener: FileWatcherEventHandler): FileWatcher {
-        return this._realFS.createFileSystemWatcher(paths, listener);
-    }
-
-    createReadStream(path: string): fs.ReadStream {
-        return this._realFS.createReadStream(this._getPartialStubOriginalPath(path));
-    }
-
-    createWriteStream(path: string): fs.WriteStream {
-        return this._realFS.createWriteStream(this._getPartialStubOriginalPath(path));
-    }
-
-    copyFileSync(src: string, dst: string): void {
-        this._realFS.copyFileSync(this._getPartialStubOriginalPath(src), this._getPartialStubOriginalPath(dst));
-    }
-
-    // Async I/O
-    readFile(path: string): Promise<Buffer> {
-        return this._realFS.readFile(this._getPartialStubOriginalPath(path));
-    }
-
-    readFileText(path: string, encoding?: BufferEncoding): Promise<string> {
-        return this._realFS.readFileText(this._getPartialStubOriginalPath(path), encoding);
-    }
-
-    // The directory returned by tmpdir must exist and be the same each time tmpdir is called.
-    tmpdir(): string {
-        return this._realFS.tmpdir();
-    }
-
-    tmpfile(options?: TmpfileOptions): string {
-        return this._realFS.tmpfile(options);
-    }
-
-    realCasePath(path: string): string {
-        return this._realFS.realCasePath(path);
-    }
-
-    getUri(originalPath: string): string {
+    override getUri(originalPath: string): string {
         const entry = this._customUriMap.get(this.getMappedFilePath(originalPath));
         if (entry) {
             return entry.uri;
@@ -310,7 +209,7 @@ export class PyrightFileSystem implements FileSystem {
                                     // not the other way around. what that means is that even if a user opens
                                     // the file explicitly from real file system (ex, vscode open file), it won't
                                     // go to the fake tmp py file. but open as it is.
-                                    this._recordVirtualFile(tmpPyFile, originalPyiFile, /* reversible */ false);
+                                    this._recordMovedEntry(tmpPyFile, originalPyiFile, /* reversible */ false);
 
                                     // This should be the only way to get to the tmp py file. and used internally
                                     // to get some info like doc string of compiled module.
@@ -319,7 +218,7 @@ export class PyrightFileSystem implements FileSystem {
                                 continue;
                             }
 
-                            this._recordVirtualFile(mappedPyiFile, originalPyiFile);
+                            this._recordMovedEntry(mappedPyiFile, originalPyiFile);
                         }
                     } catch {
                         // ignore
@@ -330,8 +229,7 @@ export class PyrightFileSystem implements FileSystem {
     }
 
     clearPartialStubs(): void {
-        this._fileMap.clear();
-        this._folderMap.clear();
+        super._clear();
 
         this._rootSearched.clear();
         this._partialStubPackagePaths.clear();
@@ -339,56 +237,12 @@ export class PyrightFileSystem implements FileSystem {
         this._conflictMap.clear();
     }
 
-    // See whether the file is mapped to another location.
-    isMappedFilePath(filepath: string): boolean {
-        return this._fileMap.has(filepath) || this._realFS.isMappedFilePath(filepath);
-    }
-
-    // Get original filepath if the given filepath is mapped.
-    getOriginalFilePath(mappedFilePath: string) {
-        return this._realFS.getOriginalFilePath(this._getPartialStubOriginalPath(mappedFilePath));
-    }
-
-    // Get mapped filepath if the given filepath is mapped.
-    getMappedFilePath(originalFilepath: string) {
-        const mappedFilePath = this._realFS.getMappedFilePath(originalFilepath);
-        return this._reverseFileMap.get(mappedFilePath) ?? mappedFilePath;
-    }
-
     // If we have a conflict file from the partial stub packages for the given file path,
     // return it.
     getConflictedFile(filepath: string) {
         return this._conflictMap.get(filepath);
     }
 
-    isInZipOrEgg(path: string): boolean {
-        return this._realFS.isInZipOrEgg(path);
-    }
-
-    private _recordVirtualFile(mappedFile: string, originalFile: string, reversible = true) {
-        this._fileMap.set(mappedFile, originalFile);
-
-        if (reversible) {
-            this._reverseFileMap.set(originalFile, mappedFile);
-        }
-
-        const directory = ensureTrailingDirectorySeparator(getDirectoryPath(mappedFile));
-        let folderInfo = this._folderMap.get(directory);
-        if (!folderInfo) {
-            folderInfo = [];
-            this._folderMap.set(directory, folderInfo);
-        }
-
-        const fileName = getFileName(mappedFile);
-        if (!folderInfo.some((entry) => entry === fileName)) {
-            folderInfo.push(fileName);
-        }
-    }
-
-    private _getPartialStubOriginalPath(mappedFilePath: string) {
-        return this._fileMap.get(mappedFilePath) ?? mappedFilePath;
-    }
-
     private _getRelativePathPartialStubs(path: string) {
         const paths: string[] = [];
 
@@ -424,7 +278,7 @@ export class PyrightFileSystem implements FileSystem {
         return paths;
     }
 
-    private _isVirtualEntry(path: string) {
-        return this._partialStubPackagePaths.has(path) || this._reverseFileMap.has(path);
+    protected override _isMovedEntry(path: string) {
+        return this._partialStubPackagePaths.has(path) || super._isMovedEntry(path);
     }
 }

packages/pyright-internal/src/readonlyAugmentedFileSystem.ts

@@ -0,0 +1,201 @@
+/*
+ * readonlyAugmentedFileSystem.ts
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ *
+ * A file system that lets one to augment backing file system but not allow
+ * modifying the backing file system.
+ */
+
+import type * as fs from 'fs';
+
+import {
+    FileSystem,
+    FileWatcher,
+    FileWatcherEventHandler,
+    MkDirOptions,
+    Stats,
+    TmpfileOptions,
+    VirtualDirent,
+} from './common/fileSystem';
+import { combinePaths, ensureTrailingDirectorySeparator, getDirectoryPath, getFileName } from './common/pathUtils';
+
+export class ReadOnlyAugmentedFileSystem implements FileSystem {
+    // Mapped file to original file map
+    private readonly _fileMap = new Map<string, string>();
+
+    // Original file to mapped file map
+    private readonly _reverseFileMap = new Map<string, string>();
+
+    // Mapped files per a containing folder map
+    private readonly _folderMap = new Map<string, string[]>();
+
+    constructor(protected _realFS: FileSystem) {}
+
+    existsSync(path: string): boolean {
+        if (this._isMovedEntry(path)) {
+            // Pretend partial stub folder and its files not exist
+            return false;
+        }
+
+        return this._realFS.existsSync(this._getOriginalPath(path));
+    }
+
+    mkdirSync(path: string, options?: MkDirOptions): void {
+        throw new Error('Operation is not allowed.');
+    }
+
+    chdir(path: string): void {
+        throw new Error('Operation is not allowed.');
+    }
+
+    readdirEntriesSync(path: string): fs.Dirent[] {
+        const entries = this._realFS.readdirEntriesSync(path).filter((item) => {
+            // Filter out the stub package directory.
+            return !this._isMovedEntry(combinePaths(path, item.name));
+        });
+
+        const partialStubs = this._folderMap.get(ensureTrailingDirectorySeparator(path));
+        if (!partialStubs) {
+            return entries;
+        }
+
+        return entries.concat(partialStubs.map((f) => new VirtualDirent(f, /* file */ true)));
+    }
+
+    readdirSync(path: string): string[] {
+        const entries = this._realFS.readdirSync(path).filter((item) => {
+            // Filter out the stub package directory.
+            return !this._isMovedEntry(combinePaths(path, item));
+        });
+
+        const partialStubs = this._folderMap.get(ensureTrailingDirectorySeparator(path));
+        if (!partialStubs) {
+            return entries;
+        }
+
+        return entries.concat(partialStubs);
+    }
+
+    readFileSync(path: string, encoding?: null): Buffer;
+    readFileSync(path: string, encoding: BufferEncoding): string;
+    readFileSync(path: string, encoding?: BufferEncoding | null): string | Buffer {
+        return this._realFS.readFileSync(this._getOriginalPath(path), encoding);
+    }
+
+    writeFileSync(path: string, data: string | Buffer, encoding: BufferEncoding | null): void {
+        throw new Error('Operation is not allowed.');
+    }
+
+    statSync(path: string): Stats {
+        return this._realFS.statSync(this._getOriginalPath(path));
+    }
+
+    unlinkSync(path: string): void {
+        throw new Error('Operation is not allowed.');
+    }
+
+    realpathSync(path: string): string {
+        return this._realFS.realpathSync(path);
+    }
+
+    getModulePath(): string {
+        return this._realFS.getModulePath();
+    }
+
+    createFileSystemWatcher(paths: string[], listener: FileWatcherEventHandler): FileWatcher {
+        return this._realFS.createFileSystemWatcher(paths, listener);
+    }
+
+    createReadStream(path: string): fs.ReadStream {
+        return this._realFS.createReadStream(this._getOriginalPath(path));
+    }
+
+    createWriteStream(path: string): fs.WriteStream {
+        throw new Error('Operation is not allowed.');
+    }
+
+    copyFileSync(src: string, dst: string): void {
+        throw new Error('Operation is not allowed.');
+    }
+
+    // Async I/O
+    readFile(path: string): Promise<Buffer> {
+        return this._realFS.readFile(this._getOriginalPath(path));
+    }
+
+    readFileText(path: string, encoding?: BufferEncoding): Promise<string> {
+        return this._realFS.readFileText(this._getOriginalPath(path), encoding);
+    }
+
+    // The directory returned by tmpdir must exist and be the same each time tmpdir is called.
+    tmpdir(): string {
+        return this._realFS.tmpdir();
+    }
+
+    tmpfile(options?: TmpfileOptions): string {
+        return this._realFS.tmpfile(options);
+    }
+
+    realCasePath(path: string): string {
+        return this._realFS.realCasePath(path);
+    }
+
+    getUri(originalPath: string): string {
+        return this._realFS.getUri(originalPath);
+    }
+
+    // See whether the file is mapped to another location.
+    isMappedFilePath(filepath: string): boolean {
+        return this._fileMap.has(filepath) || this._realFS.isMappedFilePath(filepath);
+    }
+
+    // Get original filepath if the given filepath is mapped.
+    getOriginalFilePath(mappedFilePath: string) {
+        return this._realFS.getOriginalFilePath(this._getOriginalPath(mappedFilePath));
+    }
+
+    // Get mapped filepath if the given filepath is mapped.
+    getMappedFilePath(originalFilepath: string) {
+        const mappedFilePath = this._realFS.getMappedFilePath(originalFilepath);
+        return this._reverseFileMap.get(mappedFilePath) ?? mappedFilePath;
+    }
+
+    isInZipOrEgg(path: string): boolean {
+        return this._realFS.isInZipOrEgg(path);
+    }
+
+    protected _recordMovedEntry(mappedFile: string, originalFile: string, reversible = true) {
+        this._fileMap.set(mappedFile, originalFile);
+
+        if (reversible) {
+            this._reverseFileMap.set(originalFile, mappedFile);
+        }
+
+        const directory = ensureTrailingDirectorySeparator(getDirectoryPath(mappedFile));
+        let folderInfo = this._folderMap.get(directory);
+        if (!folderInfo) {
+            folderInfo = [];
+            this._folderMap.set(directory, folderInfo);
+        }
+
+        const fileName = getFileName(mappedFile);
+        if (!folderInfo.some((entry) => entry === fileName)) {
+            folderInfo.push(fileName);
+        }
+    }
+
+    protected _getOriginalPath(mappedFilePath: string) {
+        return this._fileMap.get(mappedFilePath) ?? mappedFilePath;
+    }
+
+    protected _isMovedEntry(path: string) {
+        return this._reverseFileMap.has(path);
+    }
+
+    protected _clear() {
+        this._fileMap.clear();
+        this._reverseFileMap.clear();
+        this._folderMap.clear();
+    }
+}

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

@@ -0,0 +1,450 @@
+/*
+ * documentSymbolCollector.test.ts
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ *
+ * Tests documentSymbolCollector
+ */
+
+import assert from 'assert';
+import { CancellationToken } from 'vscode-languageserver';
+
+import { findNodeByOffset } from '../analyzer/parseTreeUtils';
+import { Program } from '../analyzer/program';
+import { createMapFromItems } from '../common/collectionUtils';
+import { ConfigOptions } from '../common/configOptions';
+import { TextRange } from '../common/textRange';
+import { DocumentSymbolCollector } from '../languageService/documentSymbolCollector';
+import { NameNode } from '../parser/parseNodes';
+import { Range } from './harness/fourslash/fourSlashTypes';
+import { parseAndGetTestState } from './harness/fourslash/testState';
+
+test('folder reference', () => {
+    const code = `
+// @filename: common/__init__.py
+//// from [|io2|] import tools as tools
+//// from [|io2|].tools import pathUtils as pathUtils
+
+// @filename: io2/empty.py
+//// # empty
+
+// @filename: io2/tools/__init__.py
+//// def combine(a, b):
+////     pass
+
+// @filename: io2/tools/pathUtils.py
+//// def getFilename(path):
+////     pass
+
+// @filename: test1.py
+//// from common import *
+////
+//// tools.combine(1, 1)
+//// pathUtils.getFilename("c")
+
+// @filename: test2.py
+//// from .[|io2|] import tools as t
+////
+//// t.combine(1, 1)
+
+// @filename: test3.py
+//// from .[|io2|].tools import pathUtils as p
+////
+//// p.getFilename("c")
+
+// @filename: test4.py
+//// from common import tools, pathUtils
+////
+//// tools.combine(1, 1)
+//// pathUtils.getFilename("c")
+
+// @filename: test5.py
+//// from [|io2|] import tools as tools
+//// from [|io2|].tools  import pathUtils as pathUtils
+////
+//// tools.combine(1, 1)
+//// pathUtils.getFilename("c")
+    `;
+
+    const state = parseAndGetTestState(code).state;
+
+    const ranges = state.getRangesByText().get('io2')!;
+    for (const range of ranges) {
+        verifyReferencesAtPosition(state.program, state.configOptions, 'io2', range.fileName, range.pos, ranges);
+    }
+});
+
+test('__init__ wildcard import', () => {
+    const code = `
+// @filename: common/__init__.py
+//// from io2 import [|tools|] as [|tools|]
+//// from io2.[|tools|] import pathUtils as pathUtils
+
+// @filename: io2/empty.py
+//// # empty
+
+// @filename: io2/tools/__init__.py
+//// def combine(a, b):
+////     pass
+
+// @filename: io2/tools/pathUtils.py
+//// def getFilename(path):
+////     pass
+
+// @filename: test1.py
+//// from common import *
+////
+//// [|tools|].combine(1, 1)
+//// pathUtils.getFilename("c")
+
+// @filename: test2.py
+//// from .io2 import [|tools|] as t
+////
+//// t.combine(1, 1)
+
+// @filename: test3.py
+//// from .io2.[|tools|] import pathUtils as p
+////
+//// p.getFilename("c")
+
+// @filename: test4.py
+//// from common import [|tools|], pathUtils
+////
+//// [|tools|].combine(1, 1)
+//// pathUtils.getFilename("c")
+
+// @filename: test5.py
+//// from io2 import [|tools|] as [|tools|]
+//// from io2.[|tools|]  import pathUtils as pathUtils
+////
+//// [|tools|].combine(1, 1)
+//// pathUtils.getFilename("c")
+    `;
+
+    const state = parseAndGetTestState(code).state;
+
+    const ranges = state.getRangesByText().get('tools')!;
+    for (const range of ranges) {
+        verifyReferencesAtPosition(state.program, state.configOptions, 'tools', range.fileName, range.pos, ranges);
+    }
+});
+
+test('submodule wildcard import', () => {
+    const code = `
+// @filename: common/__init__.py
+//// from io2 import tools as tools
+//// from io2.tools import [|pathUtils|] as [|pathUtils|]
+
+// @filename: io2/empty.py
+//// # empty
+
+// @filename: io2/tools/__init__.py
+//// def combine(a, b):
+////     pass
+
+// @filename: io2/tools/pathUtils.py
+//// def getFilename(path):
+////     pass
+
+// @filename: test1.py
+//// from common import *
+////
+//// tools.combine(1, 1)
+//// [|pathUtils|].getFilename("c")
+
+// @filename: test2.py
+//// from .io2 import tools as t
+////
+//// t.combine(1, 1)
+
+// @filename: test3.py
+//// from .io2.tools import [|pathUtils|] as p
+////
+//// p.getFilename("c")
+
+// @filename: test4.py
+//// from common import tools, [|pathUtils|]
+////
+//// tools.combine(1, 1)
+//// [|pathUtils|].getFilename("c")
+
+// @filename: test5.py
+//// from io2 import tools as tools
+//// from io2.tools  import [|pathUtils|] as [|pathUtils|]
+////
+//// tools.combine(1, 1)
+//// [|pathUtils|].getFilename("c")
+    `;
+
+    const state = parseAndGetTestState(code).state;
+
+    const ranges = state.getRangesByText().get('pathUtils')!;
+    for (const range of ranges) {
+        verifyReferencesAtPosition(state.program, state.configOptions, 'pathUtils', range.fileName, range.pos, ranges);
+    }
+});
+
+test('use localName import alias', () => {
+    const code = `
+// @filename: common/__init__.py
+//// from io2 import tools as [|/*marker1*/tools|]
+//// from io2.tools import pathUtils as pathUtils
+
+// @filename: io2/empty.py
+//// # empty
+
+// @filename: io2/tools/__init__.py
+//// def combine(a, b):
+////     pass
+
+// @filename: io2/tools/pathUtils.py
+//// def getFilename(path):
+////     pass
+
+// @filename: test1.py
+//// from common import *
+////
+//// [|/*marker2*/tools|].combine(1, 1)
+//// pathUtils.getFilename("c")
+
+// @filename: test2.py
+//// from .io2 import tools as t
+////
+//// t.combine(1, 1)
+
+// @filename: test3.py
+//// from .io2.tools import pathUtils as p
+////
+//// p.getFilename("c")
+
+// @filename: test4.py
+//// from common import [|/*marker3*/tools|], pathUtils
+////
+//// [|/*marker4*/tools|].combine(1, 1)
+//// pathUtils.getFilename("c")
+
+// @filename: test5.py
+//// from io2 import tools as [|/*marker5*/tools|]
+//// from io2.tools  import pathUtils as pathUtils
+////
+//// [|/*marker6*/tools|].combine(1, 1)
+//// pathUtils.getFilename("c")
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const references = state
+        .getRangesByText()
+        .get('tools')!
+        .map((r) => ({ path: r.fileName, range: state.convertPositionRange(r) }));
+
+    state.verifyFindAllReferences({
+        marker1: { references },
+        marker2: { references },
+        marker3: { references },
+        marker4: { references },
+        marker5: { references },
+        marker6: { references },
+    });
+});
+
+test('use localName import module', () => {
+    const code = `
+// @filename: common/__init__.py
+//// from io2 import [|/*marker1*/tools|] as [|tools|]
+//// from io2.[|/*marker2*/tools|] import pathUtils as pathUtils
+
+// @filename: io2/empty.py
+//// # empty
+
+// @filename: io2/tools/__init__.py
+//// def combine(a, b):
+////     pass
+
+// @filename: io2/tools/pathUtils.py
+//// def getFilename(path):
+////     pass
+
+// @filename: test1.py
+//// from common import *
+////
+//// [|tools|].combine(1, 1)
+//// pathUtils.getFilename("c")
+
+// @filename: test2.py
+//// from .io2 import [|/*marker3*/tools|] as t
+////
+//// t.combine(1, 1)
+
+// @filename: test3.py
+//// from .io2.[|/*marker4*/tools|] import pathUtils as p
+////
+//// p.getFilename("c")
+
+// @filename: test4.py
+//// from common import [|tools|], pathUtils
+////
+//// [|tools|].combine(1, 1)
+//// pathUtils.getFilename("c")
+
+// @filename: test5.py
+//// from io2 import [|/*marker5*/tools|] as [|tools|]
+//// from io2.[|/*marker6*/tools|]  import pathUtils as pathUtils
+////
+//// [|tools|].combine(1, 1)
+//// pathUtils.getFilename("c")
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const references = state
+        .getRangesByText()
+        .get('tools')!
+        .map((r) => ({ path: r.fileName, range: state.convertPositionRange(r) }));
+
+    state.verifyFindAllReferences({
+        marker1: { references },
+        marker2: { references },
+        marker3: { references },
+        marker4: { references },
+        marker5: { references },
+        marker6: { references },
+    });
+});
+
+test('import dotted name', () => {
+    const code = `
+// @filename: nest1/__init__.py
+//// # empty
+
+// @filename: nest1/nest2/__init__.py
+//// # empty
+
+// @filename: nest1/nest2/module.py
+//// def foo():
+////     pass
+
+// @filename: test1.py
+//// import [|nest1|].[|nest2|].[|module|]
+////
+//// [|nest1|].[|nest2|].[|module|]
+
+// @filename: nest1/test2.py
+//// import [|nest1|].[|nest2|].[|module|]
+////
+//// [|nest1|].[|nest2|].[|module|]
+    `;
+
+    const state = parseAndGetTestState(code).state;
+
+    function verify(name: string) {
+        const ranges = state.getRangesByText().get(name)!;
+        for (const range of ranges) {
+            verifyReferencesAtPosition(state.program, state.configOptions, name, range.fileName, range.pos, ranges);
+        }
+    }
+
+    verify('nest1');
+    verify('nest2');
+    verify('module');
+});
+
+test('import alias', () => {
+    const code = `
+// @filename: nest/__init__.py
+//// # empty
+
+// @filename: nest/module2.py
+//// # empty
+
+// @filename: module1.py
+//// # empty
+
+// @filename: test1.py
+//// import [|/*marker1*/module1|] as [|module1|]
+
+// @filename: test2.py
+//// import nest.[|/*marker2*/module2|] as [|module2|]
+    `;
+
+    const state = parseAndGetTestState(code).state;
+
+    const marker1 = state.getMarkerByName('marker1');
+    const ranges1 = state.getRangesByText().get('module1')!;
+    verifyReferencesAtPosition(
+        state.program,
+        state.configOptions,
+        'module1',
+        marker1.fileName,
+        marker1.position,
+        ranges1
+    );
+
+    const marker2 = state.getMarkerByName('marker2');
+    const ranges2 = state.getRangesByText().get('module2')!;
+    verifyReferencesAtPosition(
+        state.program,
+        state.configOptions,
+        'module2',
+        marker2.fileName,
+        marker2.position,
+        ranges2
+    );
+});
+
+test('string in __all__', () => {
+    const code = `
+// @filename: test1.py
+//// class [|/*marker1*/A|]:
+////     pass
+////
+//// a: "[|A|]" = "A"
+////
+//// __all__ = [ "[|A|]" ]
+    `;
+
+    const state = parseAndGetTestState(code).state;
+
+    const marker1 = state.getMarkerByName('marker1');
+    const ranges1 = state.getRangesByText().get('A')!;
+    verifyReferencesAtPosition(state.program, state.configOptions, 'A', marker1.fileName, marker1.position, ranges1);
+});
+
+function verifyReferencesAtPosition(
+    program: Program,
+    configOption: ConfigOptions,
+    symbolName: string,
+    fileName: string,
+    position: number,
+    ranges: Range[]
+) {
+    const sourceFile = program.getBoundSourceFile(fileName);
+    assert(sourceFile);
+
+    const node = findNodeByOffset(sourceFile.getParseResults()!.parseTree, position);
+    const decls = DocumentSymbolCollector.getDeclarationsForNode(
+        node as NameNode,
+        program.evaluator!,
+        /*resolveLocalName*/ true,
+        CancellationToken.None,
+        program.test_createSourceMapper(configOption.findExecEnvironment(fileName))
+    );
+
+    const rangesByFile = createMapFromItems(ranges, (r) => r.fileName);
+    for (const rangeFileName of rangesByFile.keys()) {
+        const collector = new DocumentSymbolCollector(
+            symbolName,
+            decls,
+            program.evaluator!,
+            CancellationToken.None,
+            program.getBoundSourceFile(rangeFileName)!.getParseResults()!.parseTree,
+            /*treatModuleInImportAndFromImportSame*/ true
+        );
+
+        const results = collector.collect();
+        const rangesOnFile = rangesByFile.get(rangeFileName)!;
+        assert.strictEqual(results.length, rangesOnFile.length, `${rangeFileName}@${symbolName}`);
+
+        for (const result of results) {
+            assert(rangesOnFile.some((r) => r.pos === result.range.start && r.end === TextRange.getEnd(result.range)));
+        }
+    }
+}

packages/pyright-internal/src/tests/fourslash/findallreferences.module.nested.fourslash.ts

@@ -0,0 +1,67 @@
+/// <reference path="fourslash.ts" />
+
+// @filename: nested/__init__.py
+//// from .[|/*module1*/module1|] import module1Func as module1Func
+
+// @filename: nested/module1.py
+//// def module1Func():
+////     pass
+
+// @filename: test1.py
+//// import [|/*nest1*/nested|].[|/*module2*/module1|]
+//// import [|/*nest2*/nested|].[|/*module3*/module1|] as m
+////
+//// [|/*nest3*/nested|].[|/*module4*/module1|].module1Func()
+
+// @filename: test2.py
+//// from [|/*nest4*/nested|].[|/*module5*/module1|] import module1Func
+//// from .[|/*nest5*/nested|].[|/*module6*/module1|] import module1Func as f
+
+// @filename: test3.py
+//// from .[|/*nest6*/nested|] import [|/*module7*/module1|]
+//// from .[|/*nest7*/nested|] import [|/*module8*/module1|] as m
+
+// @filename: code/test4.py
+//// from ..[|/*nest8*/nested|] import [|/*module9*/module1|]
+//// from ..[|/*nest9*/nested|] import [|/*module10*/module1|] as m
+//// from ..[|/*nest10*/nested|].[|/*module11*/module1|] import module1Func
+
+{
+    const nestedReferences = helper
+        .getRangesByText()
+        .get('nested')!
+        .map((r) => {
+            return { path: r.fileName, range: helper.convertPositionRange(r) };
+        });
+
+    const moduleReferences = helper
+        .getRangesByText()
+        .get('module1')!
+        .map((r) => {
+            return { path: r.fileName, range: helper.convertPositionRange(r) };
+        });
+
+    helper.verifyFindAllReferences({
+        nest1: { references: nestedReferences },
+        nest2: { references: nestedReferences },
+        nest3: { references: nestedReferences },
+        nest4: { references: nestedReferences },
+        nest5: { references: nestedReferences },
+        nest6: { references: nestedReferences },
+        nest7: { references: nestedReferences },
+        nest8: { references: nestedReferences },
+        nest9: { references: nestedReferences },
+        nest10: { references: nestedReferences },
+        module1: { references: moduleReferences },
+        module2: { references: moduleReferences },
+        module3: { references: moduleReferences },
+        module4: { references: moduleReferences },
+        module5: { references: moduleReferences },
+        module6: { references: moduleReferences },
+        module7: { references: moduleReferences },
+        module8: { references: moduleReferences },
+        module9: { references: moduleReferences },
+        module10: { references: moduleReferences },
+        module11: { references: moduleReferences },
+    });
+}

packages/pyright-internal/src/tests/fourslash/findallreferences.modules.duplicated.fourslash.ts

@@ -0,0 +1,59 @@
+/// <reference path="fourslash.ts" />
+
+// @filename: module1.py
+//// def module1Func():
+////     pass
+
+// @filename: nest/__init__.py
+//// # empty
+
+// @filename: nest/module1.py
+//// def nestModule1Func():
+////     pass
+
+// @filename: test1.py
+//// from [|/*marker1*/nest|] import [|/*marker2*/module1|]
+////
+//// from [|/*marker3*/nest|].[|/*marker4*/module1|] import module1Func
+////
+//// import [|/*marker5*/nest|].[|/*marker6*/module1|]
+//// import [|/*marker7*/module1|]
+////
+//// [|/*marker8*/nest|].[|/*marker9*/module1|]
+
+{
+    const nestReferences = helper
+        .getRangesByText()
+        .get('nest')!
+        .map((r) => {
+            return { path: r.fileName, range: helper.convertPositionRange(r) };
+        });
+
+    const marker7 = helper.getMarkerByName('marker7');
+    const module1References = helper
+        .getRangesByText()
+        .get('module1')!
+        .filter((r) => r.marker !== marker7)
+        .map((r) => {
+            return { path: r.fileName, range: helper.convertPositionRange(r) };
+        });
+
+    helper.verifyFindAllReferences({
+        marker1: { references: nestReferences },
+        marker2: { references: module1References },
+        marker3: { references: nestReferences },
+        marker4: { references: module1References },
+        marker5: { references: nestReferences },
+        marker6: { references: module1References },
+        marker7: {
+            references: helper
+                .getRanges()
+                .filter((r) => r.marker === marker7)
+                .map((r) => {
+                    return { path: r.fileName, range: helper.convertPositionRange(r) };
+                }),
+        },
+        marker8: { references: nestReferences },
+        marker9: { references: module1References },
+    });
+}

packages/pyright-internal/src/tests/fourslash/findallreferences.modules.fourslash.ts

@@ -0,0 +1,46 @@
+/// <reference path="fourslash.ts" />
+
+// @filename: module1.py
+//// def module1Func():
+////     pass
+
+// @filename: test1.py
+//// import [|/*marker1*/module1|]
+//// import [|/*marker2*/module1|] as m
+////
+//// [|/*marker3*/module1|].module1Func()
+
+// @filename: test2.py
+//// from [|/*marker4*/module1|] import module1Func
+//// from .[|/*marker5*/module1|] import module1Func as f
+
+// @filename: test3.py
+//// from . import [|/*marker6*/module1|]
+//// from . import [|/*marker7*/module1|] as m
+
+// @filename: nested/test4.py
+//// from .. import [|/*marker8*/module1|]
+//// from .. import [|/*marker9*/module1|] as m
+//// from ..[|/*marker10*/module1|] import module1Func
+
+{
+    const references = helper
+        .getRangesByText()
+        .get('module1')!
+        .map((r) => {
+            return { path: r.fileName, range: helper.convertPositionRange(r) };
+        });
+
+    helper.verifyFindAllReferences({
+        marker1: { references },
+        marker2: { references },
+        marker3: { references },
+        marker4: { references },
+        marker5: { references },
+        marker6: { references },
+        marker7: { references },
+        marker8: { references },
+        marker9: { references },
+        marker10: { references },
+    });
+}

packages/pyright-internal/src/tests/fourslash/findallreferences.modules.shadow.fourslash.ts

@@ -0,0 +1,95 @@
+/// <reference path="fourslash.ts" />
+
+// @filename: module1.py
+//// def module1Func():
+////     pass
+
+// @filename: nest1/__init__.py
+//// # empty
+
+// @filename: nest1/module1.py
+//// def nest1Module1Func():
+////     pass
+
+// @filename: nest1/nest2/__init__.py
+//// # empty
+
+// @filename: nest1/nest2/module1.py
+//// def nest2Module1Func():
+////     pass
+
+// @filename: test1.py
+//// from [|/*nest1_1*/nest1|] import [|{| "name":"nest1_module1", "target":"nest1" |}module1|]
+//// from [|/*nest1_2*/nest1|].[|/*nest2_1*/nest2|] import [|{| "name":"nest2_module1", "target":"nest2" |}module1|]
+////
+//// import [|/*nest1_3*/nest1|]
+//// import [|/*nest1_4*/nest1|].[|/*nest2_2*/nest2|]
+//// import [|/*nest1_5*/nest1|].[|/*nest2_3*/nest2|].[|{| "name":"nest2_module2", "target":"nest2" |}module1|]
+////
+//// from [|/*nest1_6*/nest1|] import [|/*nest2_4*/nest2|]
+////
+//// [|{| "name":"module4" |}module1|]
+//// [|/*nest1_7*/nest1|]
+//// [|/*nest1_8*/nest1|].[|/*nest2_5*/nest2|]
+//// [|/*nest1_9*/nest1|].[|{| "name":"module5", "target":"none" |}module1|]
+
+{
+    const nest1References = helper
+        .getRangesByText()
+        .get('nest1')!
+        .map((r) => {
+            return { path: r.fileName, range: helper.convertPositionRange(r) };
+        });
+
+    const nest2References = helper
+        .getRangesByText()
+        .get('nest2')!
+        .map((r) => {
+            return { path: r.fileName, range: helper.convertPositionRange(r) };
+        });
+
+    const nest2ModuleReferences = helper
+        .getFilteredRanges<{ target?: string }>(
+            (m, d, t) => t === 'module1' && !!d && (!d.target || d.target === 'nest2')
+        )
+        .map((r) => {
+            return { path: r.fileName, range: helper.convertPositionRange(r) };
+        });
+
+    helper.verifyFindAllReferences({
+        nest1_1: { references: nest1References },
+        nest1_2: { references: nest1References },
+        nest1_3: { references: nest1References },
+        nest1_4: { references: nest1References },
+        nest1_5: { references: nest1References },
+        nest1_6: { references: nest1References },
+        nest1_8: { references: nest1References },
+        nest1_9: { references: nest1References },
+        nest2_1: { references: nest2References },
+        nest2_2: { references: nest2References },
+        nest2_3: { references: nest2References },
+        nest2_4: { references: nest2References },
+        nest2_5: { references: nest2References },
+        nest2_module1: { references: nest2ModuleReferences },
+        nest2_module2: { references: nest2ModuleReferences },
+        nest1_module1: {
+            references: helper
+                .getFilteredRanges<{ target?: string }>(
+                    (m, d, t) => t === 'module1' && !!d && (!d.target || d.target === 'nest1')
+                )
+                .map((r) => {
+                    return { path: r.fileName, range: helper.convertPositionRange(r) };
+                }),
+        },
+        module4: {
+            references: helper
+                .getFilteredRanges<{ target?: string }>((m, d, t) => t === 'module1' && !!d && d.target !== 'none')
+                .map((r) => {
+                    return { path: r.fileName, range: helper.convertPositionRange(r) };
+                }),
+        },
+        module5: {
+            references: [],
+        },
+    });
+}

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

@@ -138,6 +138,73 @@ declare namespace _ {
         importName: string;
     }
 
+    interface TextDocumentIdentifier {
+        uri: string;
+    }
+
+    interface OptionalVersionedTextDocumentIdentifier extends TextDocumentIdentifier {
+        version: number | null;
+    }
+
+    interface AnnotatedTextEdit extends TextEdit {
+        annotationId: string;
+    }
+
+    interface TextDocumentEdit {
+        textDocument: OptionalVersionedTextDocumentIdentifier;
+        edits: (TextEdit | AnnotatedTextEdit)[];
+    }
+
+    interface FileOptions {
+        overwrite?: boolean;
+        ignoreIfExists?: boolean;
+    }
+
+    interface ResourceOperation {
+        kind: string;
+        annotationId?: string;
+    }
+
+    interface CreateFile extends ResourceOperation {
+        kind: 'create';
+        uri: string;
+        options?: FileOptions;
+    }
+
+    interface RenameFile extends ResourceOperation {
+        kind: 'rename';
+        oldUri: string;
+        newUri: string;
+        options?: FileOptions;
+    }
+
+    interface DeleteFileOptions {
+        recursive?: boolean;
+        ignoreIfNotExists?: boolean;
+    }
+
+    interface DeleteFile extends ResourceOperation {
+        kind: 'delete';
+        uri: string;
+        options?: DeleteFileOptions;
+    }
+
+    interface ChangeAnnotation {
+        label: string;
+        needsConfirmation?: boolean;
+        description?: string;
+    }
+
+    interface WorkspaceEdit {
+        changes?: {
+            [uri: string]: TextEdit[];
+        };
+        documentChanges?: (TextDocumentEdit | CreateFile | RenameFile | DeleteFile)[];
+        changeAnnotations?: {
+            [id: string]: ChangeAnnotation;
+        };
+    }
+
     type MarkupKind = 'markdown' | 'plaintext';
 
     type DefinitionFilter = 'all' | 'preferSource' | 'preferStubs';
@@ -160,6 +227,7 @@ declare namespace _ {
         getPositionRange(markerString: string): PositionRange;
         expandPositionRange(range: PositionRange, start: number, end: number): PositionRange;
         convertPositionRange(range: Range): PositionRange;
+        convertPathToUri(path: string): string;
 
         goToBOF(): void;
         goToEOF(): void;

packages/pyright-internal/src/tests/fourslash/rename.string.fourslash.ts

@@ -0,0 +1,26 @@
+/// <reference path="fourslash.ts" />
+
+// @filename: test.py
+//// class [|/*marker*/A|]:
+////     pass
+////
+//// __all__ = ["[|A|]"]
+
+// @filename: test2.py
+//// from test import [|A|]
+////
+//// a: "[|A|]" = [|A|]()
+
+{
+    helper.verifyRename({
+        marker: {
+            newName: 'RenamedA',
+            changes: helper
+                .getRangesByText()
+                .get('A')!
+                .map((r) => {
+                    return { filePath: r.fileName, range: helper.convertPositionRange(r), replacementText: 'RenamedA' };
+                }),
+        },
+    });
+}

packages/pyright-internal/src/tests/harness/fourslash/testState.ts

@@ -11,16 +11,23 @@ import assert from 'assert';
 import * as JSONC from 'jsonc-parser';
 import Char from 'typescript-char';
 import {
+    AnnotatedTextEdit,
     CancellationToken,
+    ChangeAnnotation,
     CodeAction,
     Command,
     CompletionItem,
+    CreateFile,
+    DeleteFile,
     Diagnostic,
     DocumentHighlight,
     DocumentHighlightKind,
     ExecuteCommandParams,
     MarkupContent,
     MarkupKind,
+    OptionalVersionedTextDocumentIdentifier,
+    RenameFile,
+    TextDocumentEdit,
     TextEdit,
     WorkspaceEdit,
 } from 'vscode-languageserver';
@@ -42,6 +49,7 @@ import {
     convertPathToUri,
     getBaseFileName,
     getFileExtension,
+    getFileSpec,
     normalizePath,
     normalizeSlashes,
 } from '../../../common/pathUtils';
@@ -60,7 +68,7 @@ import { PyrightFileSystem } from '../../../pyrightFileSystem';
 import { TestAccessHost } from '../testAccessHost';
 import * as host from '../testHost';
 import { stringify } from '../utils';
-import { createFromFileSystem, distlibFolder, libFolder } from '../vfs/factory';
+import { createFromFileSystem, distlibFolder, libFolder, typeshedFolder } from '../vfs/factory';
 import * as vfs from '../vfs/filesystem';
 import { parseTestData } from './fourSlashParser';
 import {
@@ -130,7 +138,7 @@ export class TestState {
         const ignoreCase = toBoolean(testData.globalOptions[GlobalMetadataOptionNames.ignoreCase]);
 
         this._cancellationToken = new TestCancellationToken();
-        const configOptions = this._convertGlobalOptionsToConfigOptions(this.testData.globalOptions);
+        const configOptions = this._convertGlobalOptionsToConfigOptions(this.testData.globalOptions, mountPaths);
 
         const sourceFiles = [];
         const files: vfs.FileSet = {};
@@ -296,6 +304,10 @@ export class TestState {
         return this.convertOffsetsToRange(range.fileName, range.pos, range.end);
     }
 
+    convertPathToUri(path: string) {
+        return convertPathToUri(this.fs, path);
+    }
+
     goToPosition(positionOrLineAndColumn: number | Position) {
         const pos = isNumber(positionOrLineAndColumn)
             ? positionOrLineAndColumn
@@ -712,6 +724,158 @@ export class TestState {
         return commandResult;
     }
 
+    protected verifyWorkspaceEdit(expected: WorkspaceEdit, actual: WorkspaceEdit) {
+        if (actual.changes) {
+            this._verifyTextEditMap(expected.changes!, actual.changes);
+        } else {
+            assert(!expected.changes);
+        }
+
+        if (actual.documentChanges) {
+            this._verifyDocumentEdits(expected.documentChanges!, actual.documentChanges);
+        } else {
+            assert(!expected.documentChanges);
+        }
+
+        if (actual.changeAnnotations) {
+            this._verifyChangeAnnotations(expected.changeAnnotations!, actual.changeAnnotations);
+        } else {
+            assert(!expected.changeAnnotations);
+        }
+    }
+
+    private _verifyChangeAnnotations(
+        expected: { [id: string]: ChangeAnnotation },
+        actual: { [id: string]: ChangeAnnotation }
+    ) {
+        assert.strictEqual(Object.entries(expected).length, Object.entries(actual).length);
+
+        for (const key of Object.keys(expected)) {
+            const expectedAnnotation = expected[key];
+            const actualAnnotation = actual[key];
+
+            // We need to improve it to test localized strings.
+            assert.strictEqual(expectedAnnotation.label, actualAnnotation.label);
+            assert.strictEqual(expectedAnnotation.description, actualAnnotation.description);
+
+            assert.strictEqual(expectedAnnotation.needsConfirmation, actualAnnotation.needsConfirmation);
+        }
+    }
+
+    private _textDocumentAreSame(
+        expected: OptionalVersionedTextDocumentIdentifier,
+        actual: OptionalVersionedTextDocumentIdentifier
+    ) {
+        return expected.version === actual.version && expected.uri === actual.uri;
+    }
+
+    private _verifyDocumentEdits(
+        expected: (TextDocumentEdit | CreateFile | RenameFile | DeleteFile)[],
+        actual: (TextDocumentEdit | CreateFile | RenameFile | DeleteFile)[]
+    ) {
+        assert.strictEqual(expected.length, actual.length);
+
+        for (const op of expected) {
+            assert(
+                actual.some((a) => {
+                    const expectedKind = TextDocumentEdit.is(op) ? 'edit' : op.kind;
+                    const actualKind = TextDocumentEdit.is(a) ? 'edit' : a.kind;
+                    if (expectedKind !== actualKind) {
+                        return false;
+                    }
+
+                    switch (expectedKind) {
+                        case 'edit': {
+                            const expectedEdit = op as TextDocumentEdit;
+                            const actualEdit = a as TextDocumentEdit;
+
+                            if (!this._textDocumentAreSame(expectedEdit.textDocument, actualEdit.textDocument)) {
+                                return false;
+                            }
+
+                            return this._textEditsAreSame(expectedEdit.edits, actualEdit.edits);
+                        }
+                        case 'create': {
+                            const expectedOp = op as CreateFile;
+                            const actualOp = a as CreateFile;
+                            return (
+                                expectedOp.kind === actualOp.kind &&
+                                expectedOp.annotationId === actualOp.annotationId &&
+                                expectedOp.uri === actualOp.uri &&
+                                expectedOp.options?.ignoreIfExists === actualOp.options?.ignoreIfExists &&
+                                expectedOp.options?.overwrite === actualOp.options?.overwrite
+                            );
+                        }
+                        case 'rename': {
+                            const expectedOp = op as RenameFile;
+                            const actualOp = a as RenameFile;
+                            return (
+                                expectedOp.kind === actualOp.kind &&
+                                expectedOp.annotationId === actualOp.annotationId &&
+                                expectedOp.oldUri === actualOp.oldUri &&
+                                expectedOp.newUri === actualOp.newUri &&
+                                expectedOp.options?.ignoreIfExists === actualOp.options?.ignoreIfExists &&
+                                expectedOp.options?.overwrite === actualOp.options?.overwrite
+                            );
+                        }
+                        case 'delete': {
+                            const expectedOp = op as DeleteFile;
+                            const actualOp = a as DeleteFile;
+                            return (
+                                expectedOp.annotationId === actualOp.annotationId &&
+                                expectedOp.kind === actualOp.kind &&
+                                expectedOp.uri === actualOp.uri &&
+                                expectedOp.options?.ignoreIfNotExists === actualOp.options?.ignoreIfNotExists &&
+                                expectedOp.options?.recursive === actualOp.options?.recursive
+                            );
+                        }
+                        default:
+                            debug.assertNever(expectedKind);
+                    }
+                })
+            );
+        }
+    }
+
+    private _verifyTextEditMap(expected: { [uri: string]: TextEdit[] }, actual: { [uri: string]: TextEdit[] }) {
+        assert.strictEqual(Object.entries(expected).length, Object.entries(actual).length);
+
+        for (const key of Object.keys(expected)) {
+            assert(this._textEditsAreSame(expected[key], actual[key]));
+        }
+    }
+
+    private _textEditsAreSame(
+        expectedEdits: (TextEdit | AnnotatedTextEdit)[],
+        actualEdits: (TextEdit | AnnotatedTextEdit)[]
+    ) {
+        if (expectedEdits.length !== actualEdits.length) {
+            return false;
+        }
+
+        for (const edit of expectedEdits) {
+            if (actualEdits.some((a) => this._textEditAreSame(edit, a))) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
+    private _textEditAreSame(expected: TextEdit, actual: TextEdit) {
+        if (!rangesAreEqual(expected.range, actual.range)) {
+            return false;
+        }
+
+        if (expected.newText !== actual.newText) {
+            return false;
+        }
+
+        const expectedAnnotation = AnnotatedTextEdit.is(expected) ? expected.annotationId : '';
+        const actualAnnotation = AnnotatedTextEdit.is(actual) ? actual.annotationId : '';
+        return expectedAnnotation === actualAnnotation;
+    }
+
     async verifyInvokeCodeAction(
         map: {
             [marker: string]: { title: string; files?: { [filePath: string]: string }; edits?: TextEdit[] };
@@ -1241,17 +1405,19 @@ export class TestState {
         return configFileNames.some((f) => comparer(getBaseFileName(file.fileName), f) === Comparison.EqualTo);
     }
 
-    private _convertGlobalOptionsToConfigOptions(globalOptions: CompilerSettings): ConfigOptions {
+    private _convertGlobalOptionsToConfigOptions(
+        globalOptions: CompilerSettings,
+        mountPaths?: Map<string, string>
+    ): ConfigOptions {
         const srtRoot: string = GlobalMetadataOptionNames.projectRoot;
         const projectRoot = normalizeSlashes(globalOptions[srtRoot] ?? vfs.MODULE_PATH);
         const configOptions = new ConfigOptions(projectRoot);
 
         // add more global options as we need them
-
-        return this._applyTestConfigOptions(configOptions);
+        return this._applyTestConfigOptions(configOptions, mountPaths);
     }
 
-    private _applyTestConfigOptions(configOptions: ConfigOptions) {
+    private _applyTestConfigOptions(configOptions: ConfigOptions, mountPaths?: Map<string, string>) {
         // Always enable "test mode".
         configOptions.internalTestMode = true;
 
@@ -1263,6 +1429,17 @@ export class TestState {
             configOptions.stubPath = normalizePath(combinePaths(vfs.MODULE_PATH, 'typings'));
         }
 
+        configOptions.include.push(getFileSpec(configOptions.projectRoot, '.'));
+        configOptions.exclude.push(getFileSpec(configOptions.projectRoot, typeshedFolder));
+        configOptions.exclude.push(getFileSpec(configOptions.projectRoot, distlibFolder));
+        configOptions.exclude.push(getFileSpec(configOptions.projectRoot, libFolder));
+
+        if (mountPaths) {
+            for (const mountPath of mountPaths.values()) {
+                configOptions.exclude.push(getFileSpec(configOptions.projectRoot, mountPath));
+            }
+        }
+
         return configOptions;
     }
 

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

@@ -337,6 +337,108 @@ test('no empty import roots', () => {
     importResolver.getImportRoots(configOptions.getDefaultExecEnvironment()).forEach((path) => assert(path));
 });
 
+test('import side by side file root', () => {
+    const files = [
+        {
+            path: combinePaths('/', 'file1.py'),
+            content: 'def test1(): ...',
+        },
+        {
+            path: combinePaths('/', 'file2.py'),
+            content: 'def test2(): ...',
+        },
+    ];
+
+    const importResult = getImportResult(files, ['file1']);
+    assert(importResult.isImportFound);
+    assert.strictEqual(1, importResult.resolvedPaths.filter((f) => f === combinePaths('/', 'file1.py')).length);
+});
+
+test('import side by side file sub folder', () => {
+    const files = [
+        {
+            path: combinePaths('/test', 'file1.py'),
+            content: 'def test1(): ...',
+        },
+        {
+            path: combinePaths('/test', 'file2.py'),
+            content: 'def test2(): ...',
+        },
+    ];
+
+    const importResult = getImportResult(files, ['file1']);
+    assert(importResult.isImportFound);
+    assert.strictEqual(1, importResult.resolvedPaths.filter((f) => f === combinePaths('/test', 'file1.py')).length);
+});
+
+test('import side by side file sub under src folder', () => {
+    const files = [
+        {
+            path: combinePaths('/src/nested', 'file1.py'),
+            content: 'def test1(): ...',
+        },
+        {
+            path: combinePaths('/src/nested', 'file2.py'),
+            content: 'def test2(): ...',
+        },
+    ];
+
+    const importResult = getImportResult(files, ['file1']);
+    assert(importResult.isImportFound);
+    assert.strictEqual(
+        1,
+        importResult.resolvedPaths.filter((f) => f === combinePaths('/src/nested', 'file1.py')).length
+    );
+});
+
+test('import file sub under containing folder', () => {
+    const files = [
+        {
+            path: combinePaths('/src/nested', 'file1.py'),
+            content: 'def test1(): ...',
+        },
+        {
+            path: combinePaths('/src/nested/nested2', 'file2.py'),
+            content: 'def test2(): ...',
+        },
+    ];
+
+    const importResult = getImportResult(files, ['file1']);
+    assert(importResult.isImportFound);
+    assert.strictEqual(
+        1,
+        importResult.resolvedPaths.filter((f) => f === combinePaths('/src/nested', 'file1.py')).length
+    );
+});
+
+test('import side by side file sub under lib folder', () => {
+    const files = [
+        {
+            path: combinePaths('/lib/site-packages/myLib', 'file1.py'),
+            content: 'def test1(): ...',
+        },
+        {
+            path: combinePaths('/lib/site-packages/myLib', 'file2.py'),
+            content: 'def test2(): ...',
+        },
+    ];
+
+    const importResult = getImportResult(files, ['file1']);
+    assert(!importResult.isImportFound);
+});
+
+test('dont walk up the root', () => {
+    const files = [
+        {
+            path: combinePaths('/', 'file1.py'),
+            content: 'def test1(): ...',
+        },
+    ];
+
+    const importResult = getImportResult(files, ['notExist'], (c) => (c.projectRoot = ''));
+    assert(!importResult.isImportFound);
+});
+
 function getImportResult(
     files: { path: string; content: string }[],
     nameParts: string[],
@@ -348,16 +450,18 @@ function getImportResult(
             /* empty */
         });
 
-    const file = combinePaths('src', 'file.py');
-    files.push({
-        path: file,
-        content: '# not used',
-    });
-
     const fs = createFileSystem(files);
     const configOptions = new ConfigOptions(normalizeSlashes('/'));
     setup(configOptions);
 
+    const file = files.length > 0 ? files[files.length - 1].path : combinePaths('src', 'file.py');
+    if (files.length === 0) {
+        files.push({
+            path: file,
+            content: '# not used',
+        });
+    }
+
     const importResolver = new ImportResolver(fs, configOptions, new TestAccessHost(fs.getModulePath(), [libraryRoot]));
     const importResult = importResolver.resolveImport(file, configOptions.findExecEnvironment(file), {
         leadingDots: 0,

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

@@ -171,7 +171,7 @@ test('getTextEditsForAutoImportSymbolAddition', () => {
 //// from sys import [|/*marker1*/{|"r":"meta_path, "|}|]path
     `;
 
-    testAddition(code, 'marker1', { name: 'meta_path' }, 'sys', ImportType.BuiltIn);
+    testAddition(code, 'marker1', { name: 'meta_path' }, 'sys');
 });
 
 test('getTextEditsForAutoImportSymbolAddition - already exist', () => {
@@ -179,7 +179,7 @@ test('getTextEditsForAutoImportSymbolAddition - already exist', () => {
 //// from sys import path[|/*marker1*/|]
     `;
 
-    testAddition(code, 'marker1', { name: 'path' }, 'sys', ImportType.BuiltIn);
+    testAddition(code, 'marker1', { name: 'path' }, 'sys');
 });
 
 test('getTextEditsForAutoImportSymbolAddition - with alias', () => {
@@ -187,7 +187,7 @@ test('getTextEditsForAutoImportSymbolAddition - with alias', () => {
 //// from sys import path[|/*marker1*/{|"r":", path as p"|}|]
     `;
 
-    testAddition(code, 'marker1', { name: 'path', alias: 'p' }, 'sys', ImportType.BuiltIn);
+    testAddition(code, 'marker1', { name: 'path', alias: 'p' }, 'sys');
 });
 
 test('getTextEditsForAutoImportSymbolAddition - multiple names', () => {
@@ -202,8 +202,7 @@ test('getTextEditsForAutoImportSymbolAddition - multiple names', () => {
             { name: 'meta_path', alias: 'm' },
             { name: 'zoom', alias: 'z' },
         ],
-        'sys',
-        ImportType.BuiltIn
+        'sys'
     );
 });
 
@@ -219,17 +218,23 @@ test('getTextEditsForAutoImportSymbolAddition - multiple names at some spot', ()
             { name: 'meta_path', alias: 'm' },
             { name: 'noon', alias: 'n' },
         ],
-        'sys',
-        ImportType.BuiltIn
+        'sys'
     );
 });
 
+test('getTextEditsForAutoImportSymbolAddition - wildcard', () => {
+    const code = `
+//// from sys import *[|/*marker1*/|]
+    `;
+
+    testAddition(code, 'marker1', [{ name: 'path' }], 'sys');
+});
+
 function testAddition(
     code: string,
     markerName: string,
     importNameInfo: ImportNameInfo | ImportNameInfo[],
-    moduleName: string,
-    importType: ImportType
+    moduleName: string
 ) {
     const state = parseAndGetTestState(code).state;
     const marker = state.getMarkerByName(markerName)!;

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

@@ -153,6 +153,16 @@ test('getWildcardRoot2', () => {
     assert.equal(p, normalizeSlashes('/users/me'));
 });
 
+test('getWildcardRoot with root', () => {
+    const p = getWildcardRoot('/', '.');
+    assert.equal(p, normalizeSlashes('/'));
+});
+
+test('getWildcardRoot with drive letter', () => {
+    const p = getWildcardRoot('c:/', '.');
+    assert.equal(p, normalizeSlashes('c:'));
+});
+
 test('reducePathComponentsEmpty', () => {
     assert.equal(reducePathComponents([]).length, 0);
 });

packages/pyright-internal/src/tests/renameModule.fromImports.test.ts

@@ -0,0 +1,858 @@
+/*
+ * renameModule.fromImports.test.ts
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ *
+ * Tests Program.RenameModule
+ */
+
+import { combinePaths, getDirectoryPath } from '../common/pathUtils';
+import { parseAndGetTestState } from './harness/fourslash/testState';
+import { testRenameModule } from './renameModuleTestUtils';
+
+test('rename just file name', () => {
+    const code = `
+// @filename: common/__init__.py
+//// from io2 import tools
+//// from io2.tools import [|{|"r":"renamedModule"|}pathUtils|] as [|{|"r":"renamedModule"|}pathUtils|]
+
+// @filename: io2/__init__.py
+//// # empty
+
+// @filename: io2/tools/__init__.py
+//// # empty
+
+// @filename: io2/tools/empty.py
+//// # empty
+
+// @filename: io2/tools/pathUtils.py
+//// def getFilename(path):
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from common import *
+////
+//// [|{|"r":"renamedModule"|}pathUtils|].getFilename("c")
+
+// @filename: test3.py
+//// from .io2.tools import [|{|"r":"renamedModule"|}pathUtils|] as p
+////
+//// p.getFilename("c")
+
+// @filename: test4.py
+//// from common import tools, [|{|"r":"renamedModule"|}pathUtils|]
+////
+//// [|{|"r":"renamedModule"|}pathUtils|].getFilename("c")
+
+// @filename: test5.py
+//// from io2.tools import [|{|"r":""|}pathUtils as pathUtils, |]empty[|{|"r":", renamedModule as renamedModule"|}|]
+////
+//// [|{|"r":"renamedModule"|}pathUtils|].getFilename("c")
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'renamedModule.py')}`);
+});
+
+test('from module - move file to nested folder', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: module.py
+//// def getFilename(path):
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from [|module|] import getFilename
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), 'common', 'moduleRenamed.py')}`,
+        'module',
+        'common.moduleRenamed'
+    );
+});
+
+test('from module - move file to parent folder', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def getFilename(path):
+////     [|/*marker*/pass|]
+
+// @filename: test.py
+//// from [|common.module|] import getFilename
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), '..', 'moduleRenamed.py')}`,
+        'common.module',
+        'moduleRenamed'
+    );
+});
+
+test('from module - move file to sibling folder', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def getFilename(path):
+////     [|/*marker*/pass|]
+
+// @filename: common1/__init__.py
+//// # empty
+
+// @filename: test.py
+//// from [|common.module|] import getFilename
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), '..', 'common1', 'moduleRenamed.py')}`,
+        'common.module',
+        'common1.moduleRenamed'
+    );
+});
+
+test('import name - move file to nested folder', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from [|{|"r":"common.sub"|}common|] import [|{|"r":"moduleRenamed"|}module|]
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('import name - move file to parent folder', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: common/sub/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from [|{|"r":"common"|}common.sub|] import [|{|"r":"moduleRenamed"|}module|]
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+        `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), '..', 'moduleRenamed.py')}`);
+});
+
+test('import name - move file to sibling folder', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common1/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from [|{|"r":"common1"|}common|] import [|{|"r":"moduleRenamed"|}module|]
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+            `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), '..', 'common1', 'moduleRenamed.py')}`
+    );
+});
+
+test('import alias - different name', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from [|{|"r":"common.sub"|}common|] import [|{|"r":"moduleRenamed"|}module|] as m
+//// m.foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('import alias - same name', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: common/sub/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from [|{|"r":"common"|}common.sub|] import [|{|"r":"moduleRenamed"|}module|] as [|{|"r":"moduleRenamed"|}module|]
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+        `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), '..', 'moduleRenamed.py')}`);
+});
+
+test('import multiple names', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from common import [|{|"r":""|}module, |]sub[|{|"r":"!n!from common.sub import moduleRenamed"|}|]
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('import multiple names with multiple deletions - edge case', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from common import sub[|{|"r":""|}, module, module|][|{|"r":"!n!from common.sub import moduleRenamed"|}|]
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+
+// @filename: test2.py
+//// from common import [|{|"r":""|}module, |]sub[|{|"r":""|}, module|][|{|"r":"!n!from common.sub import moduleRenamed"|}|]
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+
+// @filename: test3.py
+//// [|{|"r":""|}from common import module, module[|{|"r":"!n!from common.sub import moduleRenamed"|}|]
+//// |][|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('import multiple names with alias 1', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from common import [|{|"r":""|}module as m, |]sub[|{|"r":"!n!from common.sub import moduleRenamed as m"|}|]
+//// m.foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('import multiple names with alias 2', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from common import [|{|"r":""|}module as module, |]sub[|{|"r":"!n!from common.sub import moduleRenamed as moduleRenamed"|}|]
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('import multiple names with existing from import statement', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: common/sub/existing.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from common import [|{|"r":""|}module, |]sub
+//// from common.sub import existing[|{|"r":", moduleRenamed"|}|]
+////
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('import multiple names with existing from import statement with multiple deletion - edge case', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: common/sub/existing.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from common import [|{|"r":""|}module, module, |]sub
+//// from common.sub import existing[|{|"r":", moduleRenamed"|}|]
+////
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+
+// @filename: test2.py
+//// [|{|"r":""|}from common import module, module
+//// |]from common.sub import existing[|{|"r":", moduleRenamed"|}|]
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('import multiple names with existing from import statement with alias 1', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: common/sub/existing.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from common import [|{|"r":""|}module as m, |]sub
+//// from common.sub import existing[|{|"r":", moduleRenamed as m"|}|]
+////
+//// m.foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('import multiple names with existing from import statement with alias 2', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: common/sub/existing.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from common import [|{|"r":""|}module as module, |]sub
+//// from common.sub import existing[|{|"r":", moduleRenamed as moduleRenamed"|}|]
+////
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('from module multiple import names', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def getFilename(path):
+////     [|/*marker*/pass|]
+////
+//// def foo():
+////     pass
+
+// @filename: test.py
+//// from [|common.module|] import getFilename, foo
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), '..', 'moduleRenamed.py')}`,
+        'common.module',
+        'moduleRenamed'
+    );
+});
+
+test('from module relative path - same folder', () => {
+    const code = `
+// @filename: module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test.py
+//// from . import [|{|"r":"moduleRenamed"|}module|]
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'moduleRenamed.py')}`);
+});
+
+test('from module relative path - nested folder', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test.py
+//// from [|{|"r":".common"|}.|] import [|{|"r":"moduleRenamed"|}module|]
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'common', 'moduleRenamed.py')}`);
+});
+
+test('from module relative path - parent folder', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test.py
+//// from [|{|"r":"."|}.common|] import [|{|"r":"moduleRenamed"|}module|]
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), '..', 'moduleRenamed.py')}`);
+});
+
+test('from module relative path - sibling folder', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common1/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test.py
+//// from [|{|"r":".common1"|}.common|] import [|{|"r":"moduleRenamed"|}module|]
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), '..', 'common1', 'moduleRenamed.py')}`
+    );
+});
+
+test('from module relative path - more complex', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: base/nested/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: base/nested/test.py
+//// from [|{|"r":"...common.sub"|}...common|] import [|{|"r":"moduleRenamed"|}module|]
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('from module relative path with multiple import names', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: base/nested/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: base/nested/test.py
+//// [|{|"r":"from ...common.sub import moduleRenamed!n!"|}|]from ...common import [|{|"r":""|}module, |]sub
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('from module relative path with multiple import names and alias 1', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: base/nested/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: base/nested/test.py
+//// [|{|"r":"from ...common.sub import moduleRenamed as m!n!"|}|]from ...common import [|{|"r":""|}module as m, |]sub
+//// m.foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('from module relative path with multiple import names and alias 2', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: base/nested/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: base/nested/test.py
+//// [|{|"r":"from ...common.sub import moduleRenamed as moduleRenamed!n!"|}|]from ...common import [|{|"r":""|}module as module, |]sub
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('from module relative path with merging with existing import', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: common/sub/existing.py
+//// # empty
+
+// @filename: base/nested/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: base/nested/test.py
+//// from ...common import [|{|"r":""|}module, |]sub
+//// from ...common.sub import existing[|{|"r":", moduleRenamed"|}|]
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('from module relative path with merging with existing import with alias 1', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: common/sub/existing.py
+//// # empty
+
+// @filename: base/nested/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: base/nested/test.py
+//// from ...common import [|{|"r":""|}module as m, |]sub
+//// from ...common.sub import existing[|{|"r":", moduleRenamed as m"|}|]
+//// m.foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('from module relative path with merging with existing import with alias 2', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// # empty
+
+// @filename: common/sub/existing.py
+//// # empty
+
+// @filename: base/nested/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: base/nested/test.py
+//// from ...common import [|{|"r":""|}module as module, |]sub
+//// from ...common.sub import existing[|{|"r":", moduleRenamed as moduleRenamed"|}|]
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});
+
+test('from import move to current folder', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// def getFilename(path):
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from [|{|"r":"."|}common|] import ([|{|"r":"renamedModule"|}module|])
+////
+//// [|{|"r":"renamedModule"|}module|].getFilename("c")
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), '..', 'renamedModule.py')}`);
+});
+
+test('re-exported symbols', () => {
+    const code = `
+// @filename: common/__init__.py
+//// from [|{|"r":"common"|}common.io.nest|] import [|{|"r":"renamedModule"|}module|] as [|{|"r":"renamedModule"|}module|]
+
+// @filename: common/io/__init__.py
+//// from [|{|"r":".."|}.nest|] import [|{|"r":"renamedModule"|}module|] as [|{|"r":"renamedModule"|}module|]
+
+// @filename: common/io/nest/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: reexport.py
+//// from common import [|{|"r":"renamedModule"|}module|]
+//// __all__ = ["[|{|"r":"renamedModule"|}module|]"]
+
+// @filename: test1.py
+//// from common import [|{|"r":"renamedModule"|}module|]
+//// [|{|"r":"renamedModule"|}module|].foo()
+
+// @filename: test2.py
+//// from common.io import [|{|"r":"renamedModule"|}module|]
+//// [|{|"r":"renamedModule"|}module|].foo()
+
+// @filename: test3.py
+//// from reexport import [|{|"r":"renamedModule"|}module|]
+//// [|{|"r":"renamedModule"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), '..', '..', 'renamedModule.py')}`);
+});
+
+test('new import with existing import with wildcard', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/sub/__init__.py
+//// class A: ...
+//// __all__ = ["A"]
+
+// @filename: common/module.py
+//// def foo():
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from common import [|{|"r":""|}module, |]sub
+//// from common.sub import *[|{|"r":"!n!from common.sub import moduleRenamed"|}|]
+////
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'sub', 'moduleRenamed.py')}`);
+});

packages/pyright-internal/src/tests/renameModule.imports.test.ts

@@ -0,0 +1,450 @@
+/*
+ * renameModule.imports.test.ts
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ *
+ * Tests Program.RenameModule
+ */
+
+import { combinePaths, getDirectoryPath } from '../common/pathUtils';
+import { parseAndGetTestState } from './harness/fourslash/testState';
+import { testRenameModule } from './renameModuleTestUtils';
+
+test('rename just file name', () => {
+    const code = `
+// @filename: empty.py
+//// # empty
+
+// @filename: pathUtils.py
+//// def getFilename(path):
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// import [|pathUtils|] as p
+////
+//// p.getFilename("c")
+
+// @filename: test2.py
+//// import [|pathUtils|]
+////
+//// [|pathUtils|].getFilename("c")
+
+// @filename: test3.py
+//// import [|pathUtils|] as [|pathUtils|], empty
+////
+//// [|pathUtils|].getFilename("c")
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), 'renamedModule.py')}`,
+        'pathUtils',
+        'renamedModule'
+    );
+});
+
+test('import - move file to nested folder', () => {
+    const code = `
+// @filename: common/__init__.py
+//// def foo():
+////     pass
+
+// @filename: module.py
+//// [|/*marker*/|]
+//// # empty
+
+// @filename: test.py
+//// import [|{|"r":"common.moduleRenamed as moduleRenamed"|}module|]
+////
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'common', 'moduleRenamed.py')}`);
+});
+
+test('import - move file to parent folder', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// [|/*marker*/|]
+//// def foo():
+////     pass
+
+// @filename: test.py
+//// import [|common.module|]
+////
+//// [|common.module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), '..', 'moduleRenamed.py')}`,
+        'common.module',
+        'moduleRenamed'
+    );
+});
+
+test('import - move file to sibling folder', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: common/module.py
+//// [|/*marker*/|]
+//// def foo():
+////     pass
+
+// @filename: common1/__init__.py
+//// # empty
+
+// @filename: test.py
+//// import [|common.module|]
+////
+//// [|common.module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), '..', 'common1', 'moduleRenamed.py')}`,
+        'common.module',
+        'common1.moduleRenamed'
+    );
+});
+
+test('import alias move up file', () => {
+    const code = `
+// @filename: common/__init__.py
+//// def foo():
+////     pass
+
+// @filename: module.py
+//// [|/*marker*/|]
+//// # empty
+
+// @filename: test.py
+//// import [|{|"r":"common.moduleRenamed"|}module|] as [|{|"r":"moduleRenamed"|}module|]
+////
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+
+// @filename: test1.py
+//// import [|{|"r":"common.moduleRenamed"|}module|] as m
+////
+//// m.foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), 'common', 'moduleRenamed.py')}`,
+        'module'
+    );
+});
+
+test('import alias move down file', () => {
+    const code = `
+// @filename: common/__init__.py
+//// def foo():
+////     pass
+
+// @filename: common/module.py
+//// [|/*marker*/|]
+//// # empty
+
+// @filename: test.py
+//// import [|{|"r":"moduleRenamed"|}common.module|] as [|{|"r":"moduleRenamed"|}module|]
+////
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+
+// @filename: test1.py
+//// import [|{|"r":"moduleRenamed"|}common.module|] as m
+////
+//// m.foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), '..', 'moduleRenamed.py')}`);
+});
+
+test('import alias rename file', () => {
+    const code = `
+// @filename: module.py
+//// [|/*marker*/|]
+//// # empty
+
+// @filename: test.py
+//// import [|{|"r":"moduleRenamed"|}module|] as [|{|"r":"moduleRenamed"|}module|]
+////
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+
+// @filename: test1.py
+//// import [|{|"r":"moduleRenamed"|}module|] as m
+////
+//// m.foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'moduleRenamed.py')}`, 'module');
+});
+
+test('import alias move sibling file', () => {
+    const code = `
+// @filename: common1/__init__.py
+//// def foo():
+////     pass
+
+// @filename: common2/__init__.py
+//// def foo():
+////     pass
+
+// @filename: common1/module.py
+//// [|/*marker*/|]
+//// # empty
+
+// @filename: test.py
+//// import [|{|"r":"common2.moduleRenamed"|}common1.module|] as [|{|"r":"moduleRenamed"|}module|]
+////
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+
+// @filename: test1.py
+//// import [|{|"r":"common2.moduleRenamed"|}common1.module|] as m
+////
+//// m.foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), '..', 'common2', 'moduleRenamed.py')}`
+    );
+});
+
+test('re-export import alias through __all__', () => {
+    const code = `
+// @filename: common1/__init__.py
+//// import [|{|"r":"common2.moduleRenamed as moduleRenamed"|}module|]
+//// __all__ = ["[|{|"r":"moduleRenamed"|}module|]"]
+
+// @filename: common2/__init__.py
+//// def foo():
+////     pass
+
+// @filename: module.py
+//// [|/*marker*/|]
+//// # empty
+
+// @filename: test.py
+//// from common1 import [|{|"r":"moduleRenamed"|}module|] as [|{|"r":"moduleRenamed"|}module|]
+////
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+
+// @filename: test1.py
+//// import [|{|"r":"common2.moduleRenamed"|}module|] as m
+////
+//// m.foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'common2', 'moduleRenamed.py')}`);
+});
+
+test('re-export import alias', () => {
+    const code = `
+// @filename: common1/__init__.py
+//// import [|{|"r":"common2.moduleRenamed"|}module|] as [|{|"r":"moduleRenamed"|}module|]
+
+// @filename: common2/__init__.py
+//// def foo():
+////     pass
+
+// @filename: module.py
+//// [|/*marker*/|]
+//// # empty
+
+// @filename: test.py
+//// from common1 import [|{|"r":"moduleRenamed"|}module|] as [|{|"r":"moduleRenamed"|}module|]
+////
+//// [|{|"r":"moduleRenamed"|}module|].foo()
+
+// @filename: test1.py
+//// import [|{|"r":"common2.moduleRenamed"|}module|] as m
+////
+//// m.foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'common2', 'moduleRenamed.py')}`);
+});
+
+test('update module symbol exposed through call 1', () => {
+    const code = `
+// @filename: lib.py
+//// import reexport
+////
+//// def foo():
+////    return reexport
+
+// @filename: reexport.py
+//// import [|{|"r":"moduleRenamed"|}module|] as [|{|"r":"moduleRenamed"|}module|]
+
+// @filename: module.py
+//// [|/*marker*/|]
+//// def foo():
+////    pass
+
+// @filename: test.py
+//// from lib import foo
+////
+//// foo().[|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'moduleRenamed.py')}`);
+});
+
+test('update module symbol exposed through call 2', () => {
+    const code = `
+// @filename: lib.py
+//// import reexport
+////
+//// def foo():
+////    return reexport
+
+// @filename: reexport.py
+//// import [|{|"r":"common.moduleRenamed"|}module|] as [|{|"r":"moduleRenamed"|}module|]
+
+// @filename: module.py
+//// [|/*marker*/|]
+//// def foo():
+////    pass
+
+// @filename: test.py
+//// from lib import foo
+////
+//// foo().[|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'common', 'moduleRenamed.py')}`);
+});
+
+test('update module symbol exposed through __all__ 1', () => {
+    const code = `
+// @filename: lib.py
+//// import reexport
+////
+//// def foo():
+////    return reexport
+
+// @filename: reexport.py
+//// import [|{|"r":"moduleRenamed"|}module|]
+//// __all__ = ["[|{|"r":"moduleRenamed"|}module|]"]
+
+// @filename: module.py
+//// [|/*marker*/|]
+//// def foo():
+////    pass
+
+// @filename: test.py
+//// from lib import foo
+////
+//// foo().[|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'moduleRenamed.py')}`);
+});
+
+test('update module symbol exposed through __all__ 2', () => {
+    const code = `
+// @filename: lib.py
+//// import reexport
+////
+//// def foo():
+////    return reexport
+
+// @filename: reexport.py
+//// import [|{|"r":"common.moduleRenamed as moduleRenamed"|}module|]
+//// __all__ = ["[|{|"r":"moduleRenamed"|}module|]"]
+
+// @filename: module.py
+//// [|/*marker*/|]
+//// def foo():
+////    pass
+
+// @filename: test.py
+//// from lib import foo
+////
+//// foo().[|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'common', 'moduleRenamed.py')}`);
+});
+
+test('update module symbol exposed through __all__ 3', () => {
+    const code = `
+// @filename: lib.py
+//// import reexport
+////
+//// def foo():
+////    return reexport
+
+// @filename: reexport.py
+//// import [|{|"r":"moduleRenamed"|}common.module|] as [|{|"r":"moduleRenamed"|}module|]
+//// __all__ = ["[|{|"r":"moduleRenamed"|}module|]"]
+
+// @filename: common/module.py
+//// [|/*marker*/|]
+//// def foo():
+////    pass
+
+// @filename: test.py
+//// from lib import foo
+////
+//// foo().[|{|"r":"moduleRenamed"|}module|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), '..', 'moduleRenamed.py')}`);
+});

packages/pyright-internal/src/tests/renameModule.misc.test.ts

@@ -0,0 +1,765 @@
+/*
+ * renameModule.misc.test.ts
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ *
+ * Tests Program.RenameModule
+ */
+
+import { CancellationToken } from 'vscode-languageserver';
+
+import { assert } from '../common/debug';
+import { combinePaths, getDirectoryPath } from '../common/pathUtils';
+import { parseAndGetTestState } from './harness/fourslash/testState';
+import { testRenameModule } from './renameModuleTestUtils';
+
+test('from import with paren', () => {
+    const code = `
+// @filename: module.py
+//// def getFilename(path):
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from . import ([|module|])
+////
+//// [|module|].getFilename("c")
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), 'renamedModule.py')}`,
+        'module',
+        'renamedModule'
+    );
+});
+
+test('from import with paren with alias', () => {
+    const code = `
+// @filename: module.py
+//// def getFilename(path):
+////     [|/*marker*/pass|]
+
+// @filename: test1.py
+//// from . import ([|module|] as [|module|])
+////
+//// [|module|].getFilename("c")
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), 'renamedModule.py')}`,
+        'module',
+        'renamedModule'
+    );
+});
+
+test('from import with paren multiple import names', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: module.py
+//// def getFilename(path):
+////     [|/*marker*/pass|]
+
+// @filename: module2.py
+//// # empty
+
+// @filename: test1.py
+//// [|{|"r":"from .common import renamedModule as renamedModule!n!"|}|]from . import ([|{|"r":""|}module as module, |]module2)
+////
+//// [|{|"r":"renamedModule"|}module|].getFilename("c")
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'common', 'renamedModule.py')}`);
+});
+
+test('rename - circular references', () => {
+    const code = `
+// @filename: module1.py
+//// from . import [|mySelf|] as [|mySelf|]
+
+// @filename: mySelf.py
+//// from module1 import *
+//// [|/*marker*/mySelf|].foo()
+//// 
+//// def foo():
+////     pass
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), 'renamedModule.py')}`,
+        'mySelf',
+        'renamedModule'
+    );
+});
+
+test('move - circular references', () => {
+    const code = `
+// @filename: module1.py
+//// from [|{|"r":".common"|}.|] import [|{|"r":"renamedModule"|}mySelf|] as [|{|"r":"renamedModule"|}mySelf|]
+
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: mySelf.py
+//// [|/*marker*/|]
+//// from module1 import *
+//// [|{|"r":"renamedModule"|}mySelf|].foo()
+//// def foo():
+////     pass
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'common', 'renamedModule.py')}`);
+});
+
+test('py and pyi file update', () => {
+    const code = `
+// @filename: module.py
+//// def getFilename(path):
+////     pass
+
+// @filename: module.pyi
+//// [|/*marker*/|]
+//// def getFilename(path): ...
+
+// @filename: test1.py
+//// from . import [|module|] as [|module|]
+////
+//// [|module|].getFilename("c")
+
+// @filename: test1.pyi
+//// from . import [|module|] as [|module|]
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), 'renamedModule.pyi')}`,
+        'module',
+        'renamedModule'
+    );
+});
+
+test('py and pyi file update from py', () => {
+    // No reference. if both py and pyi exist, then given file must point to pyi not py.
+    const code = `
+// @filename: module.py
+//// [|/*marker*/|]
+//// def getFilename(path):
+////     pass
+
+// @filename: module.pyi
+//// def getFilename(path): ...
+
+// @filename: test1.py
+//// from . import module
+////
+//// module.getFilename("c")
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'renamedModule.py')}`);
+});
+
+test('handle __all__ reference', () => {
+    const code = `
+// @filename: module.py
+//// [|/*marker*/|]
+//// def getFilename(path):
+////     pass
+
+// @filename: test1.py
+//// from . import [|module|]
+////
+//// __all__ = [ "[|module|]" ]
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), 'renamedModule.py')}`,
+        'module',
+        'renamedModule'
+    );
+});
+
+test('handle __all__ re-export', () => {
+    const code = `
+// @filename: module.py
+//// [|/*marker*/|]
+//// def foo(path):
+////     pass
+
+// @filename: common/__init__.py
+//// # empty
+
+// @filename: test1.py
+//// from [|{|"r":".common"|}.|] import [|{|"r":"renamedModule"|}module|]
+////
+//// __all__ = [ "[|{|"r":"renamedModule"|}module|]" ]
+
+// @filename: test2.py
+//// from test1 import [|{|"r":"renamedModule"|}module|]
+////
+//// [|renamedModule|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'common', 'renamedModule.py')}`);
+});
+
+test('__init__.py rename', () => {
+    const code = `
+// @filename: common/__init__.py
+//// [|/*marker*/|]
+//// def getFilename(path):
+////     pass
+
+// @filename: test1.py
+//// from [|common|] import getFilename
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), 'renamedModule.py')}`,
+        'common',
+        'common.renamedModule'
+    );
+});
+
+test('__init__.py rename import', () => {
+    const code = `
+// @filename: common/__init__.py
+//// [|/*marker*/|]
+
+// @filename: test1.py
+//// import [|common|]
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), 'renamedModule.py')}`,
+        'common',
+        'common.renamedModule as renamedModule'
+    );
+});
+
+test('__init__.py move to nested folder', () => {
+    const code = `
+// @filename: common/__init__.py
+//// [|/*marker*/|]
+//// def getFilename(path):
+////     pass
+
+// @filename: test1.py
+//// from [|common|] import getFilename
+        `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), 'nested', 'renamedModule.py')}`,
+        'common',
+        'common.nested.renamedModule'
+    );
+});
+
+test('__init__.py move to nested folder with same name', () => {
+    const code = `
+// @filename: common/__init__.py
+//// [|/*marker*/|]
+//// def getFilename(path):
+////     pass
+
+// @filename: test1.py
+//// from [|common|] import getFilename
+        `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), 'nested', '__init__.py')}`,
+        'common',
+        'common.nested'
+    );
+});
+
+test('__init__.py move to parent folder', () => {
+    const code = `
+// @filename: common/__init__.py
+//// [|/*marker*/|]
+//// def getFilename(path):
+////     pass
+
+// @filename: test1.py
+//// from [|common|] import getFilename
+        `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(
+        state,
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), '..', 'renamedModule.py')}`,
+        'common',
+        'renamedModule'
+    );
+});
+
+test('__init__.py move to parent folder with same name 1', () => {
+    const code = `
+// @filename: common/__init__.py
+//// [|/*marker*/|]
+//// def getFilename(path):
+////     pass
+
+// @filename: test1.py
+//// from [|common|] import getFilename
+        `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    const edits = state.program.renameModule(
+        fileName,
+        `${combinePaths(getDirectoryPath(fileName), '..', '__init__.py')}`,
+        CancellationToken.None
+    );
+    assert(!edits);
+});
+
+test('__init__.py with alias', () => {
+    const code = `
+// @filename: common/__init__.py
+//// [|/*marker*/|]
+//// def getFilename(path):
+////     pass
+
+// @filename: test1.py
+//// from [|{|"r":".common"|}.|] import [|{|"r":"renamedModule"|}common|] as [|{|"r":"renamedModule"|}common|]
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'renamedModule.py')}`);
+});
+
+test('__init__.py import with alias', () => {
+    const code = `
+// @filename: common/__init__.py
+//// [|/*marker*/|]
+//// def getFilename(path):
+////     pass
+
+// @filename: test1.py
+//// import [|{|"r":"common.renamedModule"|}common|] as [|{|"r":"renamedModule"|}common|]
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'renamedModule.py')}`);
+});
+
+test('__init__.py rename complex', () => {
+    const code = `
+// @filename: common/__init__.py
+//// import [|{|"r":"common.nested.renamedModule"|}common.nested.lib|] as [|{|"r":"renamedModule"|}lib|]
+//// __all__ = ["[|{|"r":"renamedModule"|}lib|]"]
+
+// @filename: reexport.py
+//// from common import [|{|"r":"renamedModule"|}lib|] as [|{|"r":"renamedModule"|}lib|]
+
+// @filename: common/nested/__init__.py
+//// # empty
+
+// @filename: common/nested/lib/__init__.py
+//// [|/*marker*/|]
+//// def foo():
+////     pass
+
+// @filename: test1.py
+//// import common
+//// common.[|{|"r":"renamedModule"|}lib|].foo()
+
+// @filename: test2.py
+//// from reexport import [|{|"r":"renamedModule"|}lib|]
+//// [|{|"r":"renamedModule"|}lib|].foo()
+
+// @filename: test3.py
+//// from common import *
+//// [|{|"r":"renamedModule"|}lib|].foo()
+
+// @filename: test4.py
+//// from reexport import *
+//// [|{|"r":"renamedModule"|}lib|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), '..', 'renamedModule.py')}`);
+});
+
+test('__init__.py moved to parent folder with same name 2', () => {
+    const code = `
+// @filename: common/__init__.py
+//// import [|{|"r":"common.nested"|}common.nested.lib|] as [|{|"r":"nested"|}lib|]
+//// __all__ = ["[|{|"r":"nested"|}lib|]"]
+
+// @filename: reexport.py
+//// from common import [|{|"r":"nested"|}lib|] as [|{|"r":"nested"|}lib|]
+
+// @filename: common/nested/lib/__init__.py
+//// [|/*marker*/|]
+//// def foo():
+////     pass
+
+// @filename: test1.py
+//// import common
+//// common.[|{|"r":"nested"|}lib|].foo()
+
+// @filename: test2.py
+//// from reexport import [|{|"r":"nested"|}lib|]
+//// [|{|"r":"nested"|}lib|].foo()
+
+// @filename: test3.py
+//// from common import *
+//// [|{|"r":"nested"|}lib|].foo()
+
+// @filename: test4.py
+//// from reexport import *
+//// [|{|"r":"nested"|}lib|].foo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), '..', '__init__.py')}`);
+});
+
+test('__init__.py changes middle of dotted name', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty [|/*marker*/|]
+//// from common.nested import lib as lib
+
+// @filename: common/nested/lib.py
+//// def libFoo():
+////    pass
+
+// @filename: common/nested/__init__.py
+//// def nestedFoo():
+////     pass
+
+// @filename: test1.py
+//// import common.nested.lib
+//// common.nested.lib.libFoo()
+
+// @filename: test2.py
+//// from common import nested
+//// nested.nestedFoo()
+
+// @filename: test3.py
+//// from [|{|"r":"common.renamedModule"|}common|] import *
+//// lib.libFoo()
+
+// @filename: test4.py
+//// from [|{|"r":"common.renamedModule"|}common|] import lib
+//// lib.libFoo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'renamedModule.py')}`);
+});
+
+test('__init__.py - split from import statement', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty [|/*marker*/|]
+//// from common.nested import lib as lib
+
+// @filename: common/nested/lib.py
+//// def libFoo():
+////    pass
+
+// @filename: common/nested/__init__.py
+//// def nestedFoo():
+////     pass
+
+// @filename: test1.py
+//// from common import nested[|{|"r":""|}, lib|][|{|"r":"!n!from common.renamedModule import lib"|}|]
+//// nested.nestedFoo()
+//// lib.libFoo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'renamedModule.py')}`);
+});
+
+test('__init__.py - split from import statement with multiple names', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty [|/*marker*/|]
+//// from common.nested import lib as lib
+//// def commonFoo():
+////     pass
+
+// @filename: common/nested/lib.py
+//// def libFoo():
+////    pass
+
+// @filename: common/nested/__init__.py
+//// def nestedFoo():
+////     pass
+
+// @filename: test1.py
+//// from common import nested[|{|"r":""|}, lib, commonFoo|][|{|"r":"!n!from common.renamedModule import commonFoo, lib"|}|]
+//// nested.nestedFoo()
+//// lib.libFoo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'renamedModule.py')}`);
+});
+
+test('__init__.py - merge from import statement with multiple names', () => {
+    const code = `
+// @filename: common/nested/__init__.py
+//// # empty [|/*marker*/|]
+//// from common.nested2 import lib as lib
+//// def commonFoo():
+////     pass
+
+// @filename: common/nested/sub.py
+//// # empty
+
+// @filename: common/empty.py
+//// # empty
+
+// @filename: common/nested2/lib.py
+//// def libFoo():
+////    pass
+
+// @filename: test1.py
+//// from common.nested import [|{|"r":""|}commonFoo, lib, |]sub
+//// from common import [|{|"r":"commonFoo, "|}|]empty[|{|"r":", lib"|}|]
+////
+//// nested.commonFoo()
+//// lib.libFoo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), '..', '__init__.py')}`);
+});
+
+test('__init__.py - split from import statement with multiple names with circular reference', () => {
+    const code = `
+// @filename: common/__init__.py
+//// # empty
+//// from common.nested import lib as lib
+//// from common.nested import [|/*marker*/{|"r":"renamedModule"|}common|] as [|{|"r":"renamedModule"|}common|]
+//// 
+//// def commonFoo():
+////     pass
+
+// @filename: common/nested/lib.py
+//// def libFoo():
+////    pass
+
+// @filename: common/nested/__init__.py
+//// from [|{|"r":".."|}...|] import [|{|"r":"renamedModule"|}common|] as [|{|"r":"renamedModule"|}common|]
+
+// @filename: test1.py
+//// from common import nested[|{|"r":""|}, lib, common|][|{|"r":"!n!from common.renamedModule import lib, renamedModule"|}|]
+//// nested.[|{|"r":"renamedModule"|}common|].commonFoo()
+//// [|{|"r":"renamedModule"|}common|].commonFoo()
+//// lib.libFoo()
+    `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'renamedModule.py')}`);
+});
+
+test('__init__.py - merge from import statement with multiple names with circular reference', () => {
+    const code = `
+// @filename: common/nested/__init__.py
+//// # empty
+//// from common.nested2 import lib as lib
+//// from common.nested2 import [|/*marker*/{|"r":"common"|}nested|] as [|{|"r":"common"|}nested|]
+////
+//// def commonFoo():
+////     pass
+
+// @filename: common/nested/sub.py
+//// # empty
+
+// @filename: common/empty.py
+//// # empty
+
+// @filename: common/nested2/__init__.py
+//// from [|{|"r":"..."|}..|] import [|{|"r":"common"|}nested|] as [|{|"r":"common"|}nested|]
+
+// @filename: common/nested2/lib.py
+//// def libFoo():
+////    pass
+
+// @filename: test1.py
+//// from common.nested import [|{|"r":""|}nested, lib, |]sub
+//// from common import [|{|"r":"common, "|}|]empty[|{|"r":", lib"|}|]
+////
+//// [|{|"r":"common"|}nested|].commonFoo()
+//// lib.libFoo()
+        `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), '..', '__init__.py')}`);
+});
+
+test('__init__.py - merge from import statement with multiple names with circular reference with only name change', () => {
+    const code = `
+// @filename: common/nested/__init__.py
+//// # empty
+//// from common.nested2 import lib as lib
+//// from common.nested2 import [|/*marker*/{|"r":"renamedModule"|}nested|] as [|{|"r":"renamedModule"|}nested|]
+////
+//// def commonFoo():
+////     pass
+
+// @filename: common/nested/sub.py
+//// # empty
+
+// @filename: common/empty.py
+//// # empty
+
+// @filename: common/nested2/__init__.py
+//// from .. import [|{|"r":"renamedModule"|}nested|] as [|{|"r":"renamedModule"|}nested|]
+
+// @filename: common/nested2/lib.py
+//// def libFoo():
+////    pass
+
+// @filename: test1.py
+//// from common.nested import [|{|"r":""|}nested, lib, |]sub[|{|"r":"!n!from common.renamedModule import lib, renamedModule"|}|]
+////
+//// [|{|"r":"renamedModule"|}nested|].commonFoo()
+//// lib.libFoo()
+        `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), '..', 'renamedModule.py')}`);
+});
+
+test('add and remove consecutive edits', () => {
+    const code = `
+// @filename: a1.py
+//// # empty [|/*marker*/|]
+
+// @filename: a3.py
+//// # empty 
+
+// @filename: test1.py
+//// from . import [|{|"r":"a2"|}a1|], a3
+        `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'a2.py')}`);
+});
+
+test('add and remove consecutive edits with alias 1', () => {
+    const code = `
+// @filename: a1.py
+//// # empty [|/*marker*/|]
+
+// @filename: a3.py
+//// # empty 
+
+// @filename: test1.py
+//// from . import [|{|"r":"a2"|}a1|] as [|{|"r":"a2"|}a1|], a3
+        `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'a2.py')}`);
+});
+
+test('add and remove consecutive edits with alias 2', () => {
+    const code = `
+// @filename: a1.py
+//// # empty [|/*marker*/|]
+
+// @filename: a3.py
+//// # empty 
+
+// @filename: test1.py
+//// from . import [|{|"r":"a2"|}a1|] as a, a3
+        `;
+
+    const state = parseAndGetTestState(code).state;
+    const fileName = state.getMarkerByName('marker').fileName;
+
+    testRenameModule(state, fileName, `${combinePaths(getDirectoryPath(fileName), 'a2.py')}`);
+});

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

@@ -0,0 +1,129 @@
+/*
+ * renameModule.fromImports.test.ts
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ *
+ * Tests Program.RenameModule
+ */
+
+import assert from 'assert';
+import { CancellationToken } from 'vscode-languageserver';
+
+import { createMapFromItems } from '../common/collectionUtils';
+import { Diagnostic } from '../common/diagnostic';
+import { DiagnosticRule } from '../common/diagnosticRules';
+import { FileEditAction } from '../common/editAction';
+import { getDirectoryPath } from '../common/pathUtils';
+import { convertRangeToTextRange } from '../common/positionUtils';
+import { rangesAreEqual, TextRange } from '../common/textRange';
+import { Range } from './harness/fourslash/fourSlashTypes';
+import { TestState } from './harness/fourslash/testState';
+
+export function testRenameModule(
+    state: TestState,
+    filePath: string,
+    newFilePath: string,
+    text?: string,
+    replacementText?: string
+) {
+    const edits = state.program.renameModule(filePath, newFilePath, CancellationToken.None);
+    assert(edits);
+
+    const ranges: Range[] = [];
+    if (text !== undefined) {
+        ranges.push(...state.getRangesByText().get(text)!);
+    } else {
+        ranges.push(...state.getRanges().filter((r) => !!r.marker?.data));
+    }
+
+    assert.strictEqual(edits.length, ranges.length);
+
+    const editsPerFileMap = createMapFromItems(edits, (e) => e.filePath);
+
+    // Make sure we don't have missing imports on the original state.
+    for (const editFileName of editsPerFileMap.keys()) {
+        const sourceFile = state.program.getBoundSourceFile(editFileName)!;
+        _verifyMissingImportsDiagnostics(sourceFile.getDiagnostics(state.configOptions));
+    }
+
+    // Verify edits
+    for (const edit of edits) {
+        assert(
+            ranges.some((r) => {
+                const data = r.marker?.data as { r: string } | undefined;
+                const expectedText = replacementText ?? data?.r ?? 'N/A';
+                const expectedRange = state.convertPositionRange(r);
+                return (
+                    r.fileName === edit.filePath &&
+                    rangesAreEqual(expectedRange, edit.range) &&
+                    expectedText.replace(/!n!/g, '\n') === edit.replacementText
+                );
+            }),
+            `can't find '${replacementText ?? edit.replacementText}'@'${edit.filePath}:(${edit.range.start.line},${
+                edit.range.start.character
+            })'`
+        );
+    }
+
+    // Apply changes
+    // First apply text changes
+    for (const [editFileName, editsPerFile] of editsPerFileMap) {
+        const result = _applyEdits(state, editFileName, editsPerFile);
+        state.testFS.writeFileSync(editFileName, result.text, 'utf8');
+
+        // Update open file content if the file is in opened state.
+        if (result.version) {
+            let openedFilePath = editFileName;
+            if (editFileName === filePath) {
+                openedFilePath = newFilePath;
+                state.program.setFileClosed(filePath);
+            }
+
+            state.program.setFileOpened(openedFilePath, result.version + 1, [{ text: result.text }]);
+        }
+    }
+
+    // Second apply filename change to disk.
+    state.testFS.mkdirpSync(getDirectoryPath(newFilePath));
+    state.testFS.renameSync(filePath, newFilePath);
+
+    // Add new file as tracked file
+    state.program.addTrackedFile(newFilePath);
+
+    // And refresh program.
+    state.importResolver.invalidateCache();
+    state.program.markAllFilesDirty(true);
+
+    // Make sure we don't have missing imports after the change.
+    for (const editFileName of editsPerFileMap.keys()) {
+        const sourceFile = state.program.getBoundSourceFile(editFileName)!;
+        _verifyMissingImportsDiagnostics(sourceFile.getDiagnostics(state.configOptions));
+    }
+}
+
+function _verifyMissingImportsDiagnostics(diagnostics: Diagnostic[] | undefined) {
+    assert(
+        !diagnostics || diagnostics.filter((d) => d.getRule() === DiagnosticRule.reportMissingImports).length === 0,
+        JSON.stringify(diagnostics!.map((d) => d.message))
+    );
+}
+
+function _applyEdits(state: TestState, filePath: string, edits: FileEditAction[]) {
+    const sourceFile = state.program.getBoundSourceFile(filePath)!;
+    const parseResults = sourceFile.getParseResults()!;
+
+    const editsWithOffset = edits
+        .map((e) => ({
+            range: convertRangeToTextRange(e.range, parseResults.tokenizerOutput.lines)!,
+            text: e.replacementText,
+        }))
+        .sort((e1, e2) => e2.range.start - e1.range.start);
+
+    // Apply change in reverse order.
+    let current = parseResults.text;
+    for (const change of editsWithOffset) {
+        current = current.substr(0, change.range.start) + change.text + current.substr(TextRange.getEnd(change.range));
+    }
+
+    return { version: sourceFile.getClientVersion(), text: current };
+}