Introduced UnboundVariable and UndefinedVariable diagnostic rules and refactored auto import and add tests (#663)

client/schemas/pyrightconfig.schema.json

@@ -326,6 +326,18 @@
       "title": "Controls reporting assert expressions that will always evaluate to true",
       "default": "warning"
     },
+    "reportUndefinedVariable": {
+      "$id": "#/properties/reportUndefinedVariable",
+      "$ref": "#/definitions/diagnostic",
+      "title": "Controls reporting of attempts to use an undefined variable",
+      "default": "error"
+    },
+    "reportUnboundVariable": {
+      "$id": "#/properties/reportUnboundVariable",
+      "$ref": "#/definitions/diagnostic",
+      "title": "Controls reporting of attempts to use an unbound or possibly unbound variable",
+      "default": "error"
+    },
     "reportImplicitStringConcatenation": {
       "$id": "#/properties/reportImplicitStringConcatenation",
       "$ref": "#/definitions/diagnostic",

docs/configuration.md

@@ -44,7 +44,7 @@ The following settings control pyright’s diagnostic output (warnings or errors
 
 **enableTypeIgnoreComments** [boolean]: PEP 484 defines support for "# type: ignore" comments. This switch enables or disables support for these comments. The default value for this setting is 'true'.
 
-**reportGeneralTypeIssues** [boolean or string, optional]: Generate or suppress diagnostics for general type inconsistencies, unbound symbols, unsupported operations, argument/parameter mismatches, etc. This covers all of the basic type-checking rules not covered by other rules. It does not include syntax errors. The default value for this setting is 'error'.
+**reportGeneralTypeIssues** [boolean or string, optional]: Generate or suppress diagnostics for general type inconsistencies, unsupported operations, argument/parameter mismatches, etc. This covers all of the basic type-checking rules not covered by other rules. It does not include syntax errors. The default value for this setting is 'error'.
 
 **reportTypeshedErrors** [boolean or string, optional]: Generate or suppress diagnostics for typeshed type stub files. In general, these type stub files should be “clean” and generate no errors. The default value for this setting is 'none'.
 
@@ -116,6 +116,10 @@ The following settings control pyright’s diagnostic output (warnings or errors
 
 **reportImplicitStringConcatenation** [boolean or string, optional]: Generate or suppress diagnostics for two or more string literals that follow each other, indicating an implicit concatenation. This is considered a bad practice and often masks bugs such as missing commas. The default value for this setting is 'none'.
 
+**reportUndefinedVariable** [boolean or string, optional]: Generate or suppress diagnostics for undefined variables. The default value for this setting is 'error'.
+
+**reportUnboundVariable** [boolean or string, optional]: Generate or suppress diagnostics for unbound and possibly unbound variables. The default value for this setting is 'error'.
+
 
 ## Execution Environment Options
 Pyright allows multiple “execution environments” to be defined for different portions of your source tree. For example, a subtree may be designed to run with different import search paths or a different version of the python interpreter than the rest of the source base.

server/package-lock.json

@@ -9977,7 +9977,8 @@
         },
         "yargs-parser": {
             "version": "13.1.1",
-            "resolved": "",
+            "resolved": "https://registry.npmjs.org/yargs-parser/-/yargs-parser-13.1.1.tgz",
+            "integrity": "sha512-oVAVsHz6uFrg3XQheFII8ESO2ssAf9luWuAd6Wexsu4F3OtIW0o8IribPXYrD4WC24LWtPrJlGy87y5udK+dxQ==",
             "dev": true,
             "requires": {
                 "camelcase": "^5.0.0",

server/src/analyzer/program.ts

@@ -33,10 +33,15 @@ import {
     normalizePath,
     stripFileExtension,
 } from '../common/pathUtils';
-import { convertPositionToOffset } from '../common/positionUtils';
+import { convertPositionToOffset, convertRangeToTextRange } from '../common/positionUtils';
 import { DocumentRange, doRangesOverlap, Position, Range } from '../common/textRange';
 import { Duration, timingStats } from '../common/timing';
-import { ModuleSymbolMap } from '../languageService/completionProvider';
+import {
+    AutoImporter,
+    AutoImportResult,
+    buildModuleSymbolsMap,
+    ModuleSymbolMap,
+} from '../languageService/autoImporter';
 import { HoverResults } from '../languageService/hoverProvider';
 import { SignatureHelpResults } from '../languageService/signatureHelpProvider';
 import { ImportLookupResult } from './analyzerFileInfo';
@@ -44,9 +49,10 @@ import * as AnalyzerNodeInfo from './analyzerNodeInfo';
 import { CircularDependency } from './circularDependency';
 import { ImportResolver } from './importResolver';
 import { ImportResult, ImportType } from './importResult';
+import { findNodeByOffset } from './parseTreeUtils';
 import { Scope } from './scope';
+import { getScopeForNode } from './scopeUtils';
 import { SourceFile } from './sourceFile';
-import { SymbolTable } from './symbol';
 import { createTypeEvaluator, TypeEvaluator } from './typeEvaluator';
 import { TypeStubWriter } from './typeStubWriter';
 
@@ -580,19 +586,11 @@ export class Program {
 
     // Build a map of all modules within this program and the module-
     // level scope that contains the symbol table for the module.
-    private _buildModuleSymbolsMap(sourceFileToExclude?: SourceFileInfo): ModuleSymbolMap {
-        const moduleSymbolMap = new Map<string, SymbolTable>();
-
-        this._sourceFileList.forEach((fileInfo) => {
-            if (fileInfo !== sourceFileToExclude) {
-                const symbolTable = fileInfo.sourceFile.getModuleSymbolTable();
-                if (symbolTable) {
-                    moduleSymbolMap.set(fileInfo.sourceFile.getFilePath(), symbolTable);
-                }
-            }
-        });
-
-        return moduleSymbolMap;
+    private _buildModuleSymbolsMap(sourceFileToExclude: SourceFileInfo, token: CancellationToken): ModuleSymbolMap {
+        return buildModuleSymbolsMap(
+            this._sourceFileList.filter((s) => s !== sourceFileToExclude).map((s) => s.sourceFile),
+            token
+        );
     }
 
     private _checkTypes(fileToCheck: SourceFileInfo) {
@@ -766,6 +764,56 @@ export class Program {
         }
     }
 
+    getAutoImports(
+        filePath: string,
+        range: Range,
+        similarityLimit: number,
+        token: CancellationToken
+    ): AutoImportResult[] {
+        const sourceFileInfo = this._sourceFileMap.get(filePath);
+        if (!sourceFileInfo) {
+            return [];
+        }
+
+        const sourceFile = sourceFileInfo.sourceFile;
+        const fileContents = sourceFile.getFileContents();
+        if (!fileContents) {
+            // this only works with opened file
+            return [];
+        }
+
+        return this._runEvaluatorWithCancellationToken(token, () => {
+            this._bindFile(sourceFileInfo);
+
+            const parseTree = sourceFile.getParseResults()!;
+            const textRange = convertRangeToTextRange(range, parseTree.tokenizerOutput.lines);
+            if (!textRange) {
+                return [];
+            }
+
+            const currentNode = findNodeByOffset(parseTree.parseTree, textRange.start);
+            if (!currentNode) {
+                return [];
+            }
+
+            const word = fileContents.substr(textRange.start, textRange.length);
+            const map = this._buildModuleSymbolsMap(sourceFileInfo, token);
+            const autoImporter = new AutoImporter(
+                this._configOptions,
+                sourceFile.getFilePath(),
+                this._importResolver,
+                parseTree,
+                map
+            );
+
+            // Filter out any name that is already defined in the current scope.
+            const currentScope = getScopeForNode(currentNode);
+            return autoImporter
+                .getAutoImportCandidates(word, similarityLimit, [], token)
+                .filter((r) => !currentScope.lookUpSymbolRecursive(r.name));
+        });
+    }
+
     getDiagnostics(options: ConfigOptions): FileDiagnostics[] {
         const fileDiagnostics: FileDiagnostics[] = this._removeUnneededFiles();
 
@@ -951,7 +999,7 @@ export class Program {
                 this._importResolver,
                 this._lookUpImport,
                 this._evaluator,
-                () => this._buildModuleSymbolsMap(sourceFileInfo),
+                () => this._buildModuleSymbolsMap(sourceFileInfo, token),
                 token
             );
         });
@@ -993,7 +1041,7 @@ export class Program {
                 this._importResolver,
                 this._lookUpImport,
                 this._evaluator,
-                () => this._buildModuleSymbolsMap(sourceFileInfo),
+                () => this._buildModuleSymbolsMap(sourceFileInfo, token),
                 completionItem,
                 token
             );

server/src/analyzer/service.ts

@@ -179,6 +179,10 @@ export class AnalyzerService {
         return this._program.getBoundSourceFile(path)?.getParseResults();
     }
 
+    getAutoImports(filePath: string, range: Range, similarityLimit: number, token: CancellationToken) {
+        return this._program.getAutoImports(filePath, range, similarityLimit, token);
+    }
+
     getDefinitionForPosition(
         filePath: string,
         position: Position,

server/src/analyzer/sourceFile.ts

@@ -28,7 +28,8 @@ import * as StringUtils from '../common/stringUtils';
 import { DocumentRange, getEmptyRange, Position, TextRange } from '../common/textRange';
 import { TextRangeCollection } from '../common/textRangeCollection';
 import { timingStats } from '../common/timing';
-import { CompletionItemData, CompletionProvider, ModuleSymbolMap } from '../languageService/completionProvider';
+import { ModuleSymbolMap } from '../languageService/autoImporter';
+import { CompletionItemData, CompletionProvider } from '../languageService/completionProvider';
 import { DefinitionProvider } from '../languageService/definitionProvider';
 import { DocumentSymbolProvider } from '../languageService/documentSymbolProvider';
 import { HoverProvider, HoverResults } from '../languageService/hoverProvider';

server/src/analyzer/typeEvaluator.ts

@@ -2569,15 +2569,15 @@ export function createTypeEvaluator(importLookup: ImportLookup): TypeEvaluator {
 
             if (isUnbound(type)) {
                 addDiagnostic(
-                    fileInfo.diagnosticRuleSet.reportGeneralTypeIssues,
-                    DiagnosticRule.reportGeneralTypeIssues,
+                    fileInfo.diagnosticRuleSet.reportUnboundVariable,
+                    DiagnosticRule.reportUnboundVariable,
                     `"${name}" is unbound`,
                     node
                 );
             } else if (isPossiblyUnbound(type)) {
                 addDiagnostic(
-                    fileInfo.diagnosticRuleSet.reportGeneralTypeIssues,
-                    DiagnosticRule.reportGeneralTypeIssues,
+                    fileInfo.diagnosticRuleSet.reportUnboundVariable,
+                    DiagnosticRule.reportUnboundVariable,
                     `"${name}" is possibly unbound`,
                     node
                 );
@@ -2588,8 +2588,8 @@ export function createTypeEvaluator(importLookup: ImportLookup): TypeEvaluator {
             // Handle the special case of "reveal_type".
             if (name !== 'reveal_type') {
                 addDiagnostic(
-                    fileInfo.diagnosticRuleSet.reportGeneralTypeIssues,
-                    DiagnosticRule.reportGeneralTypeIssues,
+                    fileInfo.diagnosticRuleSet.reportUndefinedVariable,
+                    DiagnosticRule.reportUndefinedVariable,
                     `"${name}" is not defined`,
                     node
                 );

server/src/common/configOptions.ts

@@ -168,6 +168,12 @@ export interface DiagnosticRuleSet {
 
     // Report implicit concatenation of string literals.
     reportImplicitStringConcatenation: DiagnosticLevel;
+
+    // Report usage of undefined variables.
+    reportUndefinedVariable: DiagnosticLevel;
+
+    // Report usage of unbound or possibly unbound variables.
+    reportUnboundVariable: DiagnosticLevel;
 }
 
 export function cloneDiagnosticRuleSet(diagSettings: DiagnosticRuleSet): DiagnosticRuleSet {
@@ -226,6 +232,8 @@ export function getDiagLevelDiagnosticRules() {
         DiagnosticRule.reportAssertAlwaysTrue,
         DiagnosticRule.reportSelfClsParameterName,
         DiagnosticRule.reportImplicitStringConcatenation,
+        DiagnosticRule.reportUndefinedVariable,
+        DiagnosticRule.reportUnboundVariable,
     ];
 }
 
@@ -277,6 +285,8 @@ export function getStrictDiagnosticRuleSet(): DiagnosticRuleSet {
         reportAssertAlwaysTrue: 'error',
         reportSelfClsParameterName: 'error',
         reportImplicitStringConcatenation: 'none',
+        reportUnboundVariable: 'error',
+        reportUndefinedVariable: 'error',
     };
 
     return diagSettings;
@@ -324,6 +334,8 @@ export function getNoTypeCheckingDiagnosticRuleSet(): DiagnosticRuleSet {
         reportAssertAlwaysTrue: 'none',
         reportSelfClsParameterName: 'none',
         reportImplicitStringConcatenation: 'none',
+        reportUnboundVariable: 'warning',
+        reportUndefinedVariable: 'warning',
     };
 
     return diagSettings;
@@ -371,6 +383,8 @@ export function getDefaultDiagnosticRuleSet(): DiagnosticRuleSet {
         reportAssertAlwaysTrue: 'warning',
         reportSelfClsParameterName: 'warning',
         reportImplicitStringConcatenation: 'none',
+        reportUnboundVariable: 'error',
+        reportUndefinedVariable: 'error',
     };
 
     return diagSettings;
@@ -899,6 +913,20 @@ export class ConfigOptions {
                 DiagnosticRule.reportImplicitStringConcatenation,
                 defaultSettings.reportImplicitStringConcatenation
             ),
+
+            // Read the "reportUndefinedVariable" entry.
+            reportUndefinedVariable: this._convertDiagnosticLevel(
+                configObj.reportUndefinedVariable,
+                DiagnosticRule.reportUndefinedVariable,
+                defaultSettings.reportUndefinedVariable
+            ),
+
+            // Read the "reportUnboundVariable" entry.
+            reportUnboundVariable: this._convertDiagnosticLevel(
+                configObj.reportUnboundVariable,
+                DiagnosticRule.reportUnboundVariable,
+                defaultSettings.reportUnboundVariable
+            ),
         };
 
         // Read the "venvPath".

server/src/common/diagnosticRules.ts

@@ -50,4 +50,6 @@ export const enum DiagnosticRule {
     reportAssertAlwaysTrue = 'reportAssertAlwaysTrue',
     reportSelfClsParameterName = 'reportSelfClsParameterName',
     reportImplicitStringConcatenation = 'reportImplicitStringConcatenation',
+    reportUndefinedVariable = 'reportUndefinedVariable',
+    reportUnboundVariable = 'reportUnboundVariable',
 }

server/src/common/extensions.ts

@@ -7,8 +7,13 @@
  * Extension methods to various types.
  */
 
-/* eslint-disable @typescript-eslint/no-empty-function */
+// Jest won't load index.d.ts so put it in the same file.
+declare interface Promise<T> {
+    // Catches task error and ignores them.
+    ignoreErrors(): void;
+}
 
+/* eslint-disable @typescript-eslint/no-empty-function */
 // Explicitly tells that promise should be run asynchronously.
 Promise.prototype.ignoreErrors = function <T>(this: Promise<T>) {
     this.catch(() => {});

server/src/common/fileSystem.ts

@@ -46,6 +46,7 @@ export interface FileSystem {
     existsSync(path: string): boolean;
     mkdirSync(path: string, options?: MkDirOptions | number): void;
     chdir(path: string): void;
+    readdirEntriesSync(path: string): fs.Dirent[];
     readdirSync(path: string): string[];
     readFileSync(path: string, encoding?: null): Buffer;
     readFileSync(path: string, encoding: string): string;
@@ -58,6 +59,7 @@ export interface FileSystem {
     createFileSystemWatcher(paths: string[], listener: FileWatcherEventHandler): FileWatcher;
     createReadStream(path: string): fs.ReadStream;
     createWriteStream(path: string): fs.WriteStream;
+    copyFileSync(src: string, dst: string): void;
     // Async I/O
     readFile(path: string): Promise<Buffer>;
     readFileText(path: string, encoding?: string): Promise<string>;
@@ -98,9 +100,12 @@ class RealFileSystem implements FileSystem {
         process.chdir(path);
     }
 
-    readdirSync(path: string) {
+    readdirSync(path: string): string[] {
         return fs.readdirSync(path);
     }
+    readdirEntriesSync(path: string): fs.Dirent[] {
+        return fs.readdirSync(path, { withFileTypes: true });
+    }
 
     readFileSync(path: string, encoding?: null): Buffer;
     readFileSync(path: string, encoding: string): string;
@@ -143,6 +148,10 @@ class RealFileSystem implements FileSystem {
         return fs.createWriteStream(path);
     }
 
+    copyFileSync(src: string, dst: string): void {
+        fs.copyFileSync(src, dst);
+    }
+
     readFile(path: string): Promise<Buffer> {
         const d = createDeferred<Buffer>();
         fs.readFile(path, (e, b) => {

server/src/common/positionUtils.ts

@@ -56,3 +56,17 @@ export function convertPositionToOffset(position: Position, lines: TextRangeColl
 
     return lines.getItemAt(position.line).start + position.character;
 }
+
+export function convertRangeToTextRange(range: Range, lines: TextRangeCollection<TextRange>): TextRange | undefined {
+    const start = convertPositionToOffset(range.start, lines);
+    if (!start) {
+        return undefined;
+    }
+
+    const end = convertPositionToOffset(range.end, lines);
+    if (!end) {
+        return undefined;
+    }
+
+    return TextRange.fromBounds(start, end);
+}

server/src/index.d.ts

@@ -1,13 +0,0 @@
-/*
- * index.d.ts
- * Copyright (c) Microsoft Corporation.
- * Licensed under the MIT license.
- * Author: Eric Traut
- *
- * Global definitions of extension interfaces.
- */
-
-declare interface Promise<T> {
-    // Catches task error and ignores them.
-    ignoreErrors(): void;
-}

server/src/languageService/autoImporter.ts

@@ -0,0 +1,234 @@
+/*
+ * autoImporter.ts
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ *
+ */
+
+import { CancellationToken } from 'vscode-languageserver';
+
+import { ImportResolver, ModuleNameAndType } from '../analyzer/importResolver';
+import { ImportType } from '../analyzer/importResult';
+import {
+    getImportGroup,
+    getTextEditsForAutoImportInsertion,
+    getTextEditsForAutoImportSymbolAddition,
+    getTopLevelImports,
+    ImportGroup,
+    ImportStatements,
+} from '../analyzer/importStatementUtils';
+import { SourceFile } from '../analyzer/sourceFile';
+import { SymbolTable } from '../analyzer/symbol';
+import { Symbol } from '../analyzer/symbol';
+import * as SymbolNameUtils from '../analyzer/symbolNameUtils';
+import { throwIfCancellationRequested } from '../common/cancellationUtils';
+import { ConfigOptions } from '../common/configOptions';
+import { TextEditAction } from '../common/editAction';
+import { combinePaths, getDirectoryPath, getFileName, stripFileExtension } from '../common/pathUtils';
+import * as StringUtils from '../common/stringUtils';
+import { ParseNodeType } from '../parser/parseNodes';
+import { ParseResults } from '../parser/parser';
+
+export type ModuleSymbolMap = Map<string, SymbolTable>;
+
+// Build a map of all modules within this program and the module-
+// level scope that contains the symbol table for the module.
+export function buildModuleSymbolsMap(files: SourceFile[], token: CancellationToken): ModuleSymbolMap {
+    const moduleSymbolMap = new Map<string, SymbolTable>();
+
+    files.forEach((file) => {
+        throwIfCancellationRequested(token);
+        const symbolTable = file.getModuleSymbolTable();
+        if (symbolTable) {
+            moduleSymbolMap.set(file.getFilePath(), symbolTable);
+        }
+    });
+
+    return moduleSymbolMap;
+}
+
+export interface AutoImportResult {
+    name: string;
+    symbol?: Symbol;
+    source: string;
+    edits: TextEditAction[];
+}
+
+export class AutoImporter {
+    constructor(
+        private _configOptions: ConfigOptions,
+        private _filePath: string,
+        private _importResolver: ImportResolver,
+        private _parseResults: ParseResults,
+        private _moduleSymbolMap: ModuleSymbolMap
+    ) {}
+
+    getAutoImportCandidates(word: string, similarityLimit: number, excludes: string[], token: CancellationToken) {
+        const results: AutoImportResult[] = [];
+
+        const importStatements = getTopLevelImports(this._parseResults.parseTree);
+        this._moduleSymbolMap.forEach((symbolTable, filePath) => {
+            throwIfCancellationRequested(token);
+
+            const fileName = stripFileExtension(getFileName(filePath));
+
+            // Don't offer imports from files that are named with private
+            // naming semantics like "_ast.py".
+            if (SymbolNameUtils.isPrivateOrProtectedName(fileName)) {
+                return;
+            }
+
+            symbolTable.forEach((symbol, name) => {
+                throwIfCancellationRequested(token);
+
+                // For very short matching strings, we will require an exact match. Otherwise
+                // we will tend to return a list that's too long. Once we get beyond two
+                // characters, we can do a fuzzy match.
+                const isSimilar = this._isSimilar(word, name, similarityLimit);
+                if (!isSimilar || symbol.isExternallyHidden()) {
+                    return;
+                }
+
+                const alreadyIncluded = this._containsName(name, undefined, excludes, results);
+                if (alreadyIncluded) {
+                    return;
+                }
+
+                const declarations = symbol.getDeclarations();
+                if (!declarations || declarations.length === 0) {
+                    return;
+                }
+
+                // Don't include imported symbols, only those that
+                // are declared within this file.
+                if (declarations[0].path !== filePath) {
+                    return;
+                }
+
+                let importSource: string;
+                let importGroup = ImportGroup.Local;
+                let moduleNameAndType: ModuleNameAndType | undefined;
+
+                const localImport = importStatements.mapByFilePath.get(filePath);
+                if (localImport) {
+                    importSource = localImport.moduleName;
+                    importGroup = getImportGroup(localImport);
+                } else {
+                    moduleNameAndType = this._getModuleNameAndTypeFromFilePath(filePath);
+                    importSource = moduleNameAndType.moduleName;
+                    importGroup = this._getImportGroupFromModuleNameAndType(moduleNameAndType);
+                }
+
+                const autoImportTextEdits = this._getTextEditsForAutoImportByFilePath(
+                    name,
+                    importStatements,
+                    filePath,
+                    importSource,
+                    importGroup
+                );
+
+                results.push({ name, symbol, source: importSource, edits: autoImportTextEdits });
+            });
+
+            // See if this file should be offered as an implicit import.
+            const fileDir = getDirectoryPath(filePath);
+            const initPathPy = combinePaths(fileDir, '__init__.py');
+            const initPathPyi = initPathPy + 'i';
+
+            // If the current file is in a directory that also contains an "__init__.py[i]"
+            // file, we can use that directory name as an implicit import target.
+            if (!this._moduleSymbolMap.has(initPathPy) && !this._moduleSymbolMap.has(initPathPyi)) {
+                return;
+            }
+
+            const name = getFileName(fileDir);
+            const moduleNameAndType = this._getModuleNameAndTypeFromFilePath(getDirectoryPath(fileDir));
+            const importSource = moduleNameAndType.moduleName;
+            if (!importSource) {
+                return;
+            }
+
+            const isSimilar = this._isSimilar(word, name, similarityLimit);
+            if (!isSimilar) {
+                return;
+            }
+
+            const alreadyIncluded = this._containsName(name, importSource, excludes, results);
+            if (alreadyIncluded) {
+                return;
+            }
+
+            const importGroup = this._getImportGroupFromModuleNameAndType(moduleNameAndType);
+            const autoImportTextEdits = this._getTextEditsForAutoImportByFilePath(
+                name,
+                importStatements,
+                filePath,
+                importSource,
+                importGroup
+            );
+
+            results.push({ name, symbol: undefined, source: importSource, edits: autoImportTextEdits });
+        });
+
+        return results;
+    }
+
+    private _isSimilar(word: string, name: string, similarityLimit: number) {
+        return word.length > 2
+            ? StringUtils.computeCompletionSimilarity(word, name) > similarityLimit
+            : word.length > 0 && name.startsWith(word);
+    }
+
+    private _containsName(name: string, source: string | undefined, excludes: string[], results: AutoImportResult[]) {
+        if (excludes.find((e) => e === name)) {
+            return true;
+        }
+
+        if (results.find((r) => r.name === name && r.source === source)) {
+            return true;
+        }
+
+        return false;
+    }
+
+    // Given the file path of a module that we want to import,
+    // convert to a module name that can be used in an
+    // 'import from' statement.
+    private _getModuleNameAndTypeFromFilePath(filePath: string): ModuleNameAndType {
+        const execEnvironment = this._configOptions.findExecEnvironment(this._filePath);
+        return this._importResolver.getModuleNameForImport(filePath, execEnvironment);
+    }
+
+    private _getImportGroupFromModuleNameAndType(moduleNameAndType: ModuleNameAndType): ImportGroup {
+        let importGroup = ImportGroup.Local;
+        if (moduleNameAndType.isLocalTypingsFile || moduleNameAndType.importType === ImportType.ThirdParty) {
+            importGroup = ImportGroup.ThirdParty;
+        } else if (moduleNameAndType.importType === ImportType.BuiltIn) {
+            importGroup = ImportGroup.BuiltIn;
+        }
+
+        return importGroup;
+    }
+
+    private _getTextEditsForAutoImportByFilePath(
+        symbolName: string,
+        importStatements: ImportStatements,
+        filePath: string,
+        moduleName: string,
+        importGroup: ImportGroup
+    ): TextEditAction[] {
+        // Does an 'import from' statement already exist? If so, we'll reuse it.
+        const importStatement = importStatements.mapByFilePath.get(filePath);
+        if (importStatement && importStatement.node.nodeType === ParseNodeType.ImportFrom) {
+            return getTextEditsForAutoImportSymbolAddition(symbolName, importStatement, this._parseResults);
+        }
+
+        return getTextEditsForAutoImportInsertion(
+            symbolName,
+            importStatements,
+            moduleName,
+            importGroup,
+            this._parseResults
+        );
+    }
+}

server/src/languageService/completionProvider.ts

@@ -22,16 +22,7 @@ import { ImportLookup } from '../analyzer/analyzerFileInfo';
 import * as AnalyzerNodeInfo from '../analyzer/analyzerNodeInfo';
 import { Declaration, DeclarationType } from '../analyzer/declaration';
 import { convertDocStringToMarkdown } from '../analyzer/docStringToMarkdown';
-import { ImportedModuleDescriptor, ImportResolver, ModuleNameAndType } from '../analyzer/importResolver';
-import { ImportType } from '../analyzer/importResult';
-import {
-    getImportGroup,
-    getTextEditsForAutoImportInsertion,
-    getTextEditsForAutoImportSymbolAddition,
-    getTopLevelImports,
-    ImportGroup,
-    ImportStatements,
-} from '../analyzer/importStatementUtils';
+import { ImportedModuleDescriptor, ImportResolver } from '../analyzer/importResolver';
 import * as ParseTreeUtils from '../analyzer/parseTreeUtils';
 import { Symbol, SymbolTable } from '../analyzer/symbol';
 import * as SymbolNameUtils from '../analyzer/symbolNameUtils';
@@ -42,7 +33,6 @@ import { doForSubtypes, getMembersForClass, getMembersForModule } from '../analy
 import { throwIfCancellationRequested } from '../common/cancellationUtils';
 import { ConfigOptions } from '../common/configOptions';
 import { TextEditAction } from '../common/editAction';
-import { combinePaths, getDirectoryPath, getFileName, stripFileExtension } from '../common/pathUtils';
 import { convertOffsetToPosition, convertPositionToOffset } from '../common/positionUtils';
 import * as StringUtils from '../common/stringUtils';
 import { comparePositions, Position } from '../common/textRange';
@@ -62,6 +52,7 @@ import {
     StringNode,
 } from '../parser/parseNodes';
 import { ParseResults } from '../parser/parser';
+import { AutoImporter, ModuleSymbolMap } from './autoImporter';
 
 const _keywords: string[] = [
     // Expression keywords
@@ -168,8 +159,6 @@ const similarityLimit = 0.25;
 // We'll remember this many completions in the MRU list.
 const maxRecentCompletions = 128;
 
-export type ModuleSymbolMap = Map<string, SymbolTable>;
-
 export class CompletionProvider {
     private static _mostRecentCompletions: RecentCompletionInfo[] = [];
 
@@ -832,170 +821,46 @@ export class CompletionProvider {
         }
     }
 
-    private _getImportGroupFromModuleNameAndType(moduleNameAndType: ModuleNameAndType): ImportGroup {
-        let importGroup = ImportGroup.Local;
-        if (moduleNameAndType.isLocalTypingsFile || moduleNameAndType.importType === ImportType.ThirdParty) {
-            importGroup = ImportGroup.ThirdParty;
-        } else if (moduleNameAndType.importType === ImportType.BuiltIn) {
-            importGroup = ImportGroup.BuiltIn;
-        }
-
-        return importGroup;
-    }
-
     private _getAutoImportCompletions(priorWord: string, completionList: CompletionList) {
         const moduleSymbolMap = this._moduleSymbolsCallback();
-        const importStatements = getTopLevelImports(this._parseResults.parseTree);
-
-        moduleSymbolMap.forEach((symbolTable, filePath) => {
-            const fileName = stripFileExtension(getFileName(filePath));
-
-            // Don't offer imports from files that are named with private
-            // naming semantics like "_ast.py".
-            if (!SymbolNameUtils.isPrivateOrProtectedName(fileName)) {
-                symbolTable.forEach((symbol, name) => {
-                    // For very short matching strings, we will require an exact match. Otherwise
-                    // we will tend to return a list that's too long. Once we get beyond two
-                    // characters, we can do a fuzzy match.
-                    const isSimilar =
-                        priorWord.length > 2
-                            ? StringUtils.computeCompletionSimilarity(priorWord, name) > similarityLimit
-                            : priorWord.length > 0 && name.startsWith(priorWord);
-
-                    if (isSimilar) {
-                        if (!symbol.isExternallyHidden()) {
-                            // If there's already a local completion suggestion with
-                            // this name, don't add an auto-import suggestion with
-                            // the same name.
-                            const localDuplicate = completionList.items.find(
-                                (item) => item.label === name && !item.data.autoImport
-                            );
-                            const declarations = symbol.getDeclarations();
-                            if (declarations && declarations.length > 0 && localDuplicate === undefined) {
-                                // Don't include imported symbols, only those that
-                                // are declared within this file.
-                                if (declarations[0].path === filePath) {
-                                    const localImport = importStatements.mapByFilePath.get(filePath);
-                                    let importSource: string;
-                                    let importGroup = ImportGroup.Local;
-                                    let moduleNameAndType: ModuleNameAndType | undefined;
-
-                                    if (localImport) {
-                                        importSource = localImport.moduleName;
-                                        importGroup = getImportGroup(localImport);
-                                    } else {
-                                        moduleNameAndType = this._getModuleNameAndTypeFromFilePath(filePath);
-                                        importSource = moduleNameAndType.moduleName;
-                                        importGroup = this._getImportGroupFromModuleNameAndType(moduleNameAndType);
-                                    }
-
-                                    const autoImportTextEdits = this._getTextEditsForAutoImportByFilePath(
-                                        name,
-                                        importStatements,
-                                        filePath,
-                                        importSource,
-                                        importGroup
-                                    );
-
-                                    this._addSymbol(
-                                        name,
-                                        symbol,
-                                        priorWord,
-                                        completionList,
-                                        importSource,
-                                        undefined,
-                                        autoImportTextEdits
-                                    );
-                                }
-                            }
-                        }
-                    }
-                });
-
-                // See if this file should be offered as an implicit import.
-                const fileDir = getDirectoryPath(filePath);
-                const initPathPy = combinePaths(fileDir, '__init__.py');
-                const initPathPyi = initPathPy + 'i';
-
-                // If the current file is in a directory that also contains an "__init__.py[i]"
-                // file, we can use that directory name as an implicit import target.
-                if (moduleSymbolMap.has(initPathPy) || moduleSymbolMap.has(initPathPyi)) {
-                    const name = getFileName(fileDir);
-                    const moduleNameAndType = this._getModuleNameAndTypeFromFilePath(getDirectoryPath(fileDir));
-                    if (moduleNameAndType.moduleName) {
-                        const autoImportText = `Auto-import from ${moduleNameAndType.moduleName}`;
-
-                        const isDuplicateEntry = completionList.items.find((item) => {
-                            if (item.label === name) {
-                                // Don't add if there's already a local completion suggestion.
-                                if (!item.data.autoImport) {
-                                    return true;
-                                }
-
-                                // Don't add the same auto-import suggestion twice.
-                                if (item.data && item.data.autoImport === autoImportText) {
-                                    return true;
-                                }
-                            }
-
-                            return false;
-                        });
-
-                        if (!isDuplicateEntry) {
-                            const importGroup = this._getImportGroupFromModuleNameAndType(moduleNameAndType);
-                            const autoImportTextEdits = this._getTextEditsForAutoImportByFilePath(
-                                name,
-                                importStatements,
-                                filePath,
-                                moduleNameAndType.moduleName,
-                                importGroup
-                            );
-                            this._addNameToCompletionList(
-                                name,
-                                CompletionItemKind.Module,
-                                priorWord,
-                                completionList,
-                                name,
-                                '',
-                                autoImportText,
-                                undefined,
-                                autoImportTextEdits
-                            );
-                        }
-                    }
-                }
-            }
-        });
-    }
-
-    // Given the file path of a module that we want to import,
-    // convert to a module name that can be used in an
-    // 'import from' statement.
-    private _getModuleNameAndTypeFromFilePath(filePath: string): ModuleNameAndType {
-        const execEnvironment = this._configOptions.findExecEnvironment(this._filePath);
-        return this._importResolver.getModuleNameForImport(filePath, execEnvironment);
-    }
-
-    private _getTextEditsForAutoImportByFilePath(
-        symbolName: string,
-        importStatements: ImportStatements,
-        filePath: string,
-        moduleName: string,
-        importGroup: ImportGroup
-    ): TextEditAction[] {
-        // Does an 'import from' statement already exist? If so, we'll reuse it.
-        const importStatement = importStatements.mapByFilePath.get(filePath);
-        if (importStatement && importStatement.node.nodeType === ParseNodeType.ImportFrom) {
-            return getTextEditsForAutoImportSymbolAddition(symbolName, importStatement, this._parseResults);
-        }
-
-        return getTextEditsForAutoImportInsertion(
-            symbolName,
-            importStatements,
-            moduleName,
-            importGroup,
-            this._parseResults
+        const autoImporter = new AutoImporter(
+            this._configOptions,
+            this._filePath,
+            this._importResolver,
+            this._parseResults,
+            moduleSymbolMap
         );
+
+        for (const result of autoImporter.getAutoImportCandidates(
+            priorWord,
+            similarityLimit,
+            completionList.items.filter((i) => !i.data?.autoImport).map((i) => i.label),
+            this._cancellationToken
+        )) {
+            if (result.symbol) {
+                this._addSymbol(
+                    result.name,
+                    result.symbol,
+                    priorWord,
+                    completionList,
+                    result.source,
+                    undefined,
+                    result.edits
+                );
+            } else {
+                this._addNameToCompletionList(
+                    result.name,
+                    CompletionItemKind.Module,
+                    priorWord,
+                    completionList,
+                    result.name,
+                    '',
+                    `Auto-import from ${result.source}`,
+                    undefined,
+                    result.edits
+                );
+            }
+        }
     }
 
     private _getImportFromCompletions(importFromNode: ImportFromNode, priorWord: string): CompletionList | undefined {

server/src/tests/fourslash/completions.autoimport.Lib.Found.Module.fourslash.ts

@@ -0,0 +1,27 @@
+/// <reference path="fourslash.ts" />
+// @asynctest: true
+
+// @filename: test1.py
+//// testLib[|/*marker*/|]
+
+// @filename: test2.py
+//// import testLib
+
+// @filename: testLib/__init__.pyi
+// @library: true
+//// class Test:
+////     pass
+
+helper.verifyCompletion('included', {
+    marker: {
+        completions: [
+            {
+                label: 'testLib',
+                documentation: {
+                    kind: 'markdown',
+                    value: 'Auto-import from lib.site-packages\n\n```python\ntestLib\n```\n',
+                },
+            },
+        ],
+    },
+});

server/src/tests/fourslash/completions.autoimport.Lib.Found.Type.fourslash.ts

@@ -0,0 +1,27 @@
+/// <reference path="fourslash.ts" />
+// @asynctest: true
+
+// @filename: test1.py
+//// Test[|/*marker*/|]
+
+// @filename: test2.py
+//// import testLib
+
+// @filename: testLib/__init__.pyi
+// @library: true
+//// class Test:
+////     pass
+
+helper.verifyCompletion('included', {
+    marker: {
+        completions: [
+            {
+                label: 'Test',
+                documentation: {
+                    kind: 'markdown',
+                    value: 'Auto-import from testLib\n\n',
+                },
+            },
+        ],
+    },
+});

server/src/tests/fourslash/completions.autoimport.Lib.Found.duplication.fourslash.ts

@@ -0,0 +1,41 @@
+/// <reference path="fourslash.ts" />
+// @asynctest: true
+
+// @filename: test1.py
+//// testLib[|/*marker*/|]
+
+// @filename: test2.py
+//// import testLib
+//// import testLib.test1
+//// import testLib.test2
+//// a = testLib.test1.Test1()
+//// b = testLib.test2.Test2()
+
+// @filename: testLib/__init__.pyi
+// @library: true
+//// class Test:
+////     pass
+
+// @filename: testLib/test1.pyi
+// @library: true
+//// class Test1:
+////     pass
+
+// @filename: testLib/test2.pyi
+// @library: true
+//// class Test2:
+////     pass
+
+helper.verifyCompletion('included', {
+    marker: {
+        completions: [
+            {
+                label: 'testLib',
+                documentation: {
+                    kind: 'markdown',
+                    value: 'Auto-import from lib.site-packages\n\n```python\ntestLib\n```\n',
+                },
+            },
+        ],
+    },
+});

server/src/tests/fourslash/completions.autoimport.Lib.NotFound.fourslash.ts

@@ -0,0 +1,12 @@
+/// <reference path="fourslash.ts" />
+// @asynctest: true
+
+// @filename: test1.py
+//// Test[|/*marker*/|]
+
+// @filename: testLib/__init__.pyi
+// @library: true
+//// class Test:
+////     pass
+
+helper.verifyCompletion('included', { marker: { completions: [] } });

server/src/tests/fourslash/completions.autoimport.fourslash.ts

@@ -0,0 +1,23 @@
+/// <reference path="fourslash.ts" />
+// @asynctest: true
+
+// @filename: test1.py
+//// Test[|/*marker*/|]
+
+// @filename: test2.py
+//// class Test:
+////     pass
+
+helper.verifyCompletion('included', {
+    marker: {
+        completions: [
+            {
+                label: 'Test',
+                documentation: {
+                    kind: 'markdown',
+                    value: 'Auto-import from test2\n\n',
+                },
+            },
+        ],
+    },
+});

server/src/tests/fourslash/fourslash.ts

@@ -121,9 +121,12 @@ declare namespace _ {
             [marker: string]: { codeActions: { title: string; kind: string; command: Command }[] };
         }): void;
         verifyCommand(command: Command, files: { [filePath: string]: string }): void;
-        verifyInvokeCodeAction(map: {
-            [marker: string]: { title: string; files?: { [filePath: string]: string }; edits?: TextEdit[] };
-        }): void;
+        verifyInvokeCodeAction(
+            map: {
+                [marker: string]: { title: string; files?: { [filePath: string]: string }; edits?: TextEdit[] };
+            },
+            verifyCodeActionCount?: boolean
+        ): void;
         verifyHover(map: { [marker: string]: { value: string; kind: string } }): void;
         verifyCompletion(
             verifyMode: 'exact' | 'included' | 'excluded',

server/src/tests/harness/fourslash/runner.ts

@@ -83,13 +83,7 @@ function runCode(code: string, state: TestState): void {
             ${code}
             })`;
         const f = eval(wrappedCode);
-        f(state, Consts)
-            .then(() => {
-                markDone();
-            })
-            .catch((e: any) => {
-                markDone(e);
-            });
+        f(state, Consts);
     }
 
     function runPlainCode() {

server/src/tests/harness/fourslash/testState.ts

@@ -658,9 +658,12 @@ export class TestState {
         this.markTestDone();
     }
 
-    async verifyInvokeCodeAction(map: {
-        [marker: string]: { title: string; files?: { [filePath: string]: string }; edits?: TextEdit[] };
-    }): Promise<any> {
+    async verifyInvokeCodeAction(
+        map: {
+            [marker: string]: { title: string; files?: { [filePath: string]: string }; edits?: TextEdit[] };
+        },
+        verifyCodeActionCount?: boolean
+    ): Promise<any> {
         this._analyze();
 
         for (const range of this.getRanges()) {
@@ -672,7 +675,22 @@ export class TestState {
             const ls = new TestLanguageService(this.workspace, this.console, this.fs);
 
             const codeActions = await this._getCodeActions(range);
-            for (const codeAction of codeActions.filter((c) => c.title === map[name].title)) {
+            if (verifyCodeActionCount) {
+                if (codeActions.length !== Object.keys(map).length) {
+                    this._raiseError(
+                        `doesn't contain expected result: ${stringify(map[name])}, actual: ${stringify(codeActions)}`
+                    );
+                }
+            }
+
+            const matches = codeActions.filter((c) => c.title === map[name].title);
+            if (matches.length === 0) {
+                this._raiseError(
+                    `doesn't contain expected result: ${stringify(map[name])}, actual: ${stringify(codeActions)}`
+                );
+            }
+
+            for (const codeAction of matches) {
                 const results = await this._hostSpecificFeatures.execute(
                     ls,
                     {
@@ -686,11 +704,17 @@ export class TestState {
                     const workspaceEdits = results as WorkspaceEdit;
                     for (const edits of Object.values(workspaceEdits.changes!)) {
                         for (const edit of edits) {
-                            assert(
+                            if (
                                 map[name].edits!.filter(
                                     (e) => rangesAreEqual(e.range, edit.range) && e.newText === edit.newText
-                                ).length === 1
-                            );
+                                ).length !== 1
+                            ) {
+                                this._raiseError(
+                                    `doesn't contain expected result: ${stringify(map[name])}, actual: ${stringify(
+                                        edits
+                                    )}`
+                                );
+                            }
                         }
                     }
                 }
@@ -856,6 +880,8 @@ export class TestState {
                 assert.fail('Failed to get completions');
             }
         }
+
+        this.markTestDone();
     }
 
     verifySignature(map: {

server/src/tests/harness/vfs/filesystem.ts

@@ -7,7 +7,7 @@
  */
 
 /* eslint-disable no-dupe-class-members */
-import { ReadStream, WriteStream } from 'fs';
+import { Dirent, ReadStream, WriteStream } from 'fs';
 
 import { FileSystem, FileWatcher, FileWatcherEventHandler } from '../../../common/fileSystem';
 import * as pathUtil from '../../../common/pathUtils';
@@ -549,6 +549,25 @@ export class TestFileSystem implements FileSystem {
         return Array.from(this._getLinks(node).keys());
     }
 
+    /**
+     * Read a directory. If `path` is a symbolic link, it is dereferenced.
+     *
+     * @link http://pubs.opengroup.org/onlinepubs/9699919799/functions/readdir.html
+     *
+     * NOTE: do not rename this method as it is intended to align with the same named export of the "fs" module.
+     */
+    readdirEntriesSync(path: string): Dirent[] {
+        const { node } = this._walk(this._resolve(path));
+        if (!node) {
+            throw createIOError('ENOENT');
+        }
+        if (!isDirectory(node)) {
+            throw createIOError('ENOTDIR');
+        }
+        const entries = Array.from(this._getLinks(node).entries());
+        return entries.map(([k, v]) => makeDirEnt(k, v));
+    }
+
     /**
      * Make a directory.
      *
@@ -832,6 +851,10 @@ export class TestFileSystem implements FileSystem {
         throw new Error('Not implemented in test file system.');
     }
 
+    copyFileSync(src: string, dst: string): void {
+        throw new Error('Not implemented in test file system.');
+    }
+
     /**
      * Generates a `FileSet` patch containing all the entries in this `FileSystem` that are not in `base`.
      * @param base The base file system. If not provided, this file system's `shadowRoot` is used (if present).
@@ -1713,6 +1736,20 @@ function formatPatchWorker(dirname: string, container: FileSet): string {
     return text;
 }
 
+function makeDirEnt(name: string, node: Inode): Dirent {
+    const de: Dirent = {
+        isFile: () => isFile(node),
+        isDirectory: () => isDirectory(node),
+        isBlockDevice: () => false,
+        isCharacterDevice: () => false,
+        isFIFO: () => false,
+        isSocket: () => false,
+        isSymbolicLink: () => isSymlink(node),
+        name,
+    };
+    return de;
+}
+
 class Stats {
     dev: number;
     ino: number;