Improved wildcard import logic. It now attempts to honor __all__ assignment in target module. If __all__ is not defined, all symbols except for those starting with an underscore are imported.

server/src/analyzer/binder.ts

@@ -34,7 +34,7 @@ import { ArgumentCategory, AssertNode, AssignmentExpressionNode, AssignmentNode,
     UnaryOperationNode, WhileNode, WithNode, YieldFromNode, YieldNode } from '../parser/parseNodes';
 import * as StringTokenUtils from '../parser/stringTokenUtils';
 import { KeywordType, OperatorType } from '../parser/tokenizerTypes';
-import { AnalyzerFileInfo } from './analyzerFileInfo';
+import { AnalyzerFileInfo, ImportLookupResult } from './analyzerFileInfo';
 import * as AnalyzerNodeInfo from './analyzerNodeInfo';
 import { createKeyForReference, FlowAssignment, FlowAssignmentAlias, FlowCall, FlowCondition,
     FlowFlags, FlowLabel, FlowNode, FlowPostFinally, FlowPreFinallyGate, FlowWildcardImport,
@@ -1110,12 +1110,15 @@ export class Binder extends ParseTreeWalker {
         }
 
         if (node.isWildcardImport) {
-            if (importInfo && importInfo.implicitImports) {
+            if (importInfo) {
                 const names: string[] = [];
 
                 const lookupInfo = this._fileInfo.importLookup(resolvedPath);
                 if (lookupInfo) {
-                    lookupInfo.symbolTable.forEach((symbol, name) => {
+                    const wildcardNames = this._getWildcardImportNames(lookupInfo);
+                    wildcardNames.forEach(name => {
+                        const symbol = lookupInfo.symbolTable.get(name)!;
+
                         // Don't include the ignored names in the symbol table.
                         if (!symbol.isIgnoredForProtocolMatch()) {
                             const symbol = this._bindNameToScope(this._currentScope, name);
@@ -1379,6 +1382,52 @@ export class Binder extends ParseTreeWalker {
         return false;
     }
 
+    private _getWildcardImportNames(lookupInfo: ImportLookupResult): string[] {
+        const namesToImport: string[] = [];
+
+        // Is there an __all__ statement? If so, it overrides the normal
+        // wildcard logic.
+        const allSymbol = lookupInfo.symbolTable.get('__all__');
+        if (allSymbol) {
+            const decls = allSymbol.getDeclarations();
+
+            // For now, we handle only the case where __all__ is defined
+            // through a simple assignment. Some libraries use more complex
+            // logic like __all__.extend(X) or __all__ += X. We'll punt on
+            // those for now.
+            if (decls.length === 1 && decls[0].type === DeclarationType.Variable) {
+                const firstDecl = decls[0];
+                if (firstDecl.node.parent && firstDecl.node.parent.nodeType === ParseNodeType.Assignment) {
+                    const expr = firstDecl.node.parent.rightExpression;
+                    if (expr.nodeType === ParseNodeType.List) {
+                        expr.entries.forEach(listEntryNode => {
+                            if (listEntryNode.nodeType === ParseNodeType.StringList &&
+                                    listEntryNode.strings.length === 1 &&
+                                    listEntryNode.strings[0].nodeType === ParseNodeType.String) {
+
+                                const entryName = listEntryNode.strings[0].value;
+                                if (lookupInfo.symbolTable.get(entryName)) {
+                                    namesToImport.push(entryName);
+                                }
+                            }
+                        });
+
+                        return namesToImport;
+                    }
+                }
+            }
+        }
+
+        // Import all names that don't begin with an underscore.
+        lookupInfo.symbolTable.forEach((_, name) => {
+            if (!name.startsWith('_')) {
+                namesToImport.push(name);
+            }
+        });
+
+        return namesToImport;
+    }
+
     private _walkStatementsAndReportUnreachable(statements: StatementNode[]) {
         let reportedUnreachable = false;
 

server/src/analyzer/importStatementUtils.ts

@@ -89,7 +89,9 @@ export function getTextEditsForAutoImportSymbolAddition(symbolName: string,
             }
 
             const insertionOffset = priorImport ? TextRange.getEnd(priorImport) :
-                importStatement.node.imports[0].start;
+                (importStatement.node.imports.length > 0 ?
+                    importStatement.node.imports[0].start :
+                    importStatement.node.start + importStatement.node.length);
             const insertionPosition = convertOffsetToPosition(insertionOffset, parseResults.tokenizerOutput.lines);
 
             textEditList.push({

server/src/tests/checker.test.ts

@@ -929,6 +929,16 @@ test('Import2', () => {
     validateResults(analysisResults, 2);
 });
 
+test('Import4', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['import4.py']);
+    validateResults(analysisResults, 1);
+});
+
+test('Import6', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['import6.py']);
+    validateResults(analysisResults, 2);
+});
+
 test('Overload1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['overload1.py']);
     validateResults(analysisResults, 2);

server/src/tests/samples/import3.py

@@ -0,0 +1,8 @@
+# This sample is imported by import4.py.
+
+__all__ = ['foo', '_foo', '_bar']
+
+foo = 3
+_foo = 4
+bar = 5
+_bar = 6

server/src/tests/samples/import4.py

@@ -0,0 +1,12 @@
+# This sample tests wildcard imports.
+
+from .import3 import *
+
+a = foo
+b = _foo
+
+# This should generate an error because bar isn't
+# included in the __all__ assigment.
+c = bar
+d = _bar
+

server/src/tests/samples/import5.py

@@ -0,0 +1,6 @@
+# This sample is imported by import6.py.
+
+foo = 3
+_foo = 4
+bar = 5
+_bar = 6

server/src/tests/samples/import6.py

@@ -0,0 +1,17 @@
+# This sample tests wildcard imports.
+
+from .import5 import *
+
+a = foo
+
+# This should generate an error because there is no
+# __all__ assignment, and names starting with an underscore
+# should not be imported in a wildcard.
+b = _foo
+c = bar
+
+# This should generate an error because there is no
+# __all__ assignment, and names starting with an underscore
+# should not be imported in a wildcard.
+d = _bar
+