Added support for additional idioms for __all__ manipulation, including __all__.extend() and __all__.remove(). The "extend" form also supports merging of __all__ lists from submodules.

This commit is contained in:
Eric Traut 2020-09-28 14:45:00 -07:00
parent b6cd5eedae
commit 5ca499b112
8 changed files with 248 additions and 77 deletions

View File

@ -20,6 +20,7 @@ export type ImportLookup = (filePath: string) => ImportLookupResult | undefined;
export interface ImportLookupResult {
symbolTable: SymbolTable;
dunderAllNames: string[] | undefined;
docString?: string;
}

View File

@ -55,6 +55,9 @@ interface AnalyzerNodeInfo {
// Map of expressions used within an execution scope (module,
// function or lambda) that requires code flow analysis.
codeFlowExpressions?: Map<string, string>;
// List of __all__ symbols in the module.
dunderAllNames?: string[];
}
export type ScopedNode = ModuleNode | ClassNode | FunctionNode | LambdaNode | ListComprehensionNode;
@ -140,6 +143,16 @@ export function setCodeFlowExpressions(node: ExecutionScopeNode, map: Map<string
analyzerNode.codeFlowExpressions = map;
}
export function getDunderAllNames(node: ModuleNode): string[] | undefined {
const analyzerNode = node as AnalyzerNodeInfo;
return analyzerNode.dunderAllNames;
}
export function setDunderAllNames(node: ModuleNode, names: string[] | undefined) {
const analyzerNode = node as AnalyzerNodeInfo;
analyzerNode.dunderAllNames = names;
}
export function isCodeUnreachable(node: ParseNode): boolean {
let curNode: ParseNode | undefined = node;

View File

@ -107,7 +107,6 @@ import { NameBindingType, Scope, ScopeType } from './scope';
import * as StaticExpressions from './staticExpressions';
import { indeterminateSymbolId, Symbol, SymbolFlags } from './symbol';
import { isConstantName, isPrivateOrProtectedName } from './symbolNameUtils';
import { getNamesInDunderAll } from './symbolUtils';
interface MemberAccessInfo {
classNode: ClassNode;
@ -184,6 +183,9 @@ export class Binder extends ParseTreeWalker {
// and the names they alias to.
private _typingSymbolAliases: Map<string, string> = new Map<string, string>();
// List of names statically assigned to __all__ symbol.
private _dunderAllNames: string[] | undefined;
// Flow node that is used for unreachable code.
private static _unreachableFlowNode: FlowNode = {
flags: FlowFlags.Unreachable,
@ -201,7 +203,7 @@ export class Binder extends ParseTreeWalker {
// binding the builtins module itself.
const isBuiltInModule = this._fileInfo.builtinsScope === undefined;
this._createNewScope(
const moduleScope = this._createNewScope(
isBuiltInModule ? ScopeType.Builtin : ScopeType.Module,
this._fileInfo.builtinsScope,
() => {
@ -235,6 +237,8 @@ export class Binder extends ParseTreeWalker {
// Perform all analysis that was deferred during the first pass.
this._bindDeferred();
AnalyzerNodeInfo.setDunderAllNames(node, this._dunderAllNames);
return {
moduleDocString: ParseTreeUtils.getDocString(node.statements),
};
@ -538,6 +542,56 @@ export class Binder extends ParseTreeWalker {
this.walk(node.leftExpression);
this.walkMultiple(node.arguments);
this._createCallFlowNode(node);
// Is this an manipulation of dunder all?
if (
this._currentScope.type === ScopeType.Module &&
node.leftExpression.nodeType === ParseNodeType.MemberAccess &&
node.leftExpression.leftExpression.nodeType === ParseNodeType.Name &&
node.leftExpression.leftExpression.value === '__all__'
) {
// Is this a call to "__all__.extend()"?
if (node.leftExpression.memberName.value === 'extend' && node.arguments.length === 1) {
const argExpr = node.arguments[0].valueExpression;
// Is this a call to "__all__.extend([<list>])"?
if (argExpr.nodeType === ParseNodeType.List) {
argExpr.entries.forEach((listEntryNode) => {
if (
listEntryNode.nodeType === ParseNodeType.StringList &&
listEntryNode.strings.length === 1 &&
listEntryNode.strings[0].nodeType === ParseNodeType.String
) {
this._dunderAllNames?.push(listEntryNode.strings[0].value);
}
});
} else if (
argExpr.nodeType === ParseNodeType.MemberAccess &&
argExpr.leftExpression.nodeType === ParseNodeType.Name &&
argExpr.memberName.value === '__all__'
) {
// Is this a call to "__all__.extend(<mod>.__all__)"?
const namesToAdd = this._getDunderAllNamesFromImport(argExpr.leftExpression.value);
if (namesToAdd) {
namesToAdd.forEach((name) => {
this._dunderAllNames?.push(name);
});
}
}
} else if (node.leftExpression.memberName.value === 'remove' && node.arguments.length === 1) {
// Is this a call to "__all__.remove()"?
const argExpr = node.arguments[0].valueExpression;
if (
argExpr.nodeType === ParseNodeType.StringList &&
argExpr.strings.length === 1 &&
argExpr.strings[0].nodeType === ParseNodeType.String &&
this._dunderAllNames
) {
this._dunderAllNames = this._dunderAllNames.filter((name) => name !== argExpr.strings[0].value);
}
}
}
return false;
}
@ -570,6 +624,38 @@ export class Binder extends ParseTreeWalker {
this._createAssignmentTargetFlowNodes(node.leftExpression, /* walkTargets */ true, /* unbound */ false);
// Is this an assignment to dunder all?
if (
this._currentScope.type === ScopeType.Module &&
node.leftExpression.nodeType === ParseNodeType.Name &&
node.leftExpression.value === '__all__'
) {
const expr = node.rightExpression;
this._dunderAllNames = [];
if (expr.nodeType === ParseNodeType.List) {
expr.entries.forEach((listEntryNode) => {
if (
listEntryNode.nodeType === ParseNodeType.StringList &&
listEntryNode.strings.length === 1 &&
listEntryNode.strings[0].nodeType === ParseNodeType.String
) {
this._dunderAllNames!.push(listEntryNode.strings[0].value);
}
});
} else if (expr.nodeType === ParseNodeType.Tuple) {
expr.expressions.forEach((tupleEntryNode) => {
if (
tupleEntryNode.nodeType === ParseNodeType.StringList &&
tupleEntryNode.strings.length === 1 &&
tupleEntryNode.strings[0].nodeType === ParseNodeType.String
) {
this._dunderAllNames!.push(tupleEntryNode.strings[0].value);
}
});
}
}
return false;
}
@ -621,6 +707,27 @@ export class Binder extends ParseTreeWalker {
this._bindPossibleTupleNamedTarget(node.destExpression);
this._createAssignmentTargetFlowNodes(node.destExpression, /* walkTargets */ false, /* unbound */ false);
// Is this an assignment to dunder all?
if (
this._currentScope.type === ScopeType.Module &&
node.leftExpression.nodeType === ParseNodeType.Name &&
node.leftExpression.value === '__all__'
) {
const expr = node.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
) {
this._dunderAllNames?.push(listEntryNode.strings[0].value);
}
});
}
}
return false;
}
@ -1147,11 +1254,13 @@ export class Binder extends ParseTreeWalker {
}
const symbol = this._bindNameToScope(this._currentScope, symbolName);
if (symbol && this._fileInfo.isStubFile && !node.alias) {
// PEP 484 indicates that imported symbols should not be
// considered "reexported" from a type stub file unless
// they are imported using the "as" form.
symbol.setIsExternallyHidden();
if (symbol && !node.alias) {
if (this._fileInfo.isStubFile) {
// PEP 484 indicates that imported symbols should not be
// considered "reexported" from a type stub file unless
// they are imported using the "as" form.
symbol.setIsExternallyHidden();
}
}
const importInfo = AnalyzerNodeInfo.getImportInfo(node.module);
@ -1307,11 +1416,13 @@ export class Binder extends ParseTreeWalker {
const symbol = this._bindNameToScope(this._currentScope, nameNode.value);
if (symbol) {
if (this._fileInfo.isStubFile && !importSymbolNode.alias) {
// PEP 484 indicates that imported symbols should not be
// considered "reexported" from a type stub file unless
// they are imported using the "as" form.
symbol.setIsExternallyHidden();
if (!importSymbolNode.alias) {
if (this._fileInfo.isStubFile) {
// PEP 484 indicates that imported symbols should not be
// considered "reexported" from a type stub file unless
// they are imported using the "as" form.
symbol.setIsExternallyHidden();
}
}
// Is the import referring to an implicitly-imported module?
@ -1526,6 +1637,31 @@ export class Binder extends ParseTreeWalker {
return false;
}
// Attempts to resolve the module name, import it, and return
// its __all__ symbols.
private _getDunderAllNamesFromImport(varName: string): string[] | undefined {
const varSymbol = this._currentScope.lookUpSymbol(varName);
if (!varSymbol) {
return undefined;
}
// There should be only one declaration for the variable.
const aliasDecl = varSymbol.getDeclarations().find((decl) => decl.type === DeclarationType.Alias) as
| AliasDeclaration
| undefined;
const resolvedPath = aliasDecl?.path || aliasDecl?.submoduleFallback?.path;
if (!resolvedPath) {
return undefined;
}
const lookupInfo = this._fileInfo.importLookup(resolvedPath);
if (!lookupInfo) {
return undefined;
}
return lookupInfo.dunderAllNames;
}
private _addImplicitFromImport(node: ImportFromNode, importInfo?: ImportResult) {
const symbolName = node.module.nameParts[0].value;
const symbol = this._bindNameToScope(this._currentScope, symbolName);
@ -1634,13 +1770,13 @@ export class Binder extends ParseTreeWalker {
}
private _getWildcardImportNames(lookupInfo: ImportLookupResult): string[] {
let namesToImport = getNamesInDunderAll(lookupInfo.symbolTable);
if (namesToImport) {
return namesToImport;
// If a dunder all symbol is defined, it takes precedence.
if (lookupInfo.dunderAllNames) {
return lookupInfo.dunderAllNames;
}
// Import all names that don't begin with an underscore.
namesToImport = [];
const namesToImport: string[] = [];
lookupInfo.symbolTable.forEach((symbol, name) => {
if (!name.startsWith('_') && !symbol.isIgnoredForProtocolMatch()) {
namesToImport!.push(name);
@ -2094,8 +2230,10 @@ export class Binder extends ParseTreeWalker {
}
}
if (this._fileInfo.isStubFile && isPrivateOrProtectedName(name)) {
symbol.setIsExternallyHidden();
if (isPrivateOrProtectedName(name)) {
if (this._fileInfo.isStubFile) {
symbol.setIsExternallyHidden();
}
}
if (addedSymbols) {
@ -2190,7 +2328,8 @@ export class Binder extends ParseTreeWalker {
private _createNewScope(scopeType: ScopeType, parentScope: Scope | undefined, callback: () => void) {
const prevScope = this._currentScope;
this._currentScope = new Scope(scopeType, parentScope);
const newScope = new Scope(scopeType, parentScope);
this._currentScope = newScope;
// If this scope is an execution scope, allocate a new reference map.
const isExecutionScope =
@ -2205,6 +2344,8 @@ export class Binder extends ParseTreeWalker {
this._currentExecutionScopeReferenceMap = prevReferenceMap;
this._currentScope = prevScope;
return newScope;
}
private _addInferredTypeAssignmentForVariable(

View File

@ -718,9 +718,11 @@ export class Program {
}
const docString = sourceFileInfo.sourceFile.getModuleDocString();
const parseResults = sourceFileInfo.sourceFile.getParseResults();
return {
symbolTable,
dunderAllNames: AnalyzerNodeInfo.getDunderAllNames(parseResults!.parseTree),
docString,
};
};

View File

@ -394,15 +394,20 @@ export class SourceFile {
// Keep the parse info, but reset the analysis to the beginning.
this._isCheckingNeeded = true;
// If the file contains a wildcard import, we need to rebind
// also because the dependent import may have changed.
if (this._parseResults && this._parseResults.containsWildcardImport) {
this._parseTreeNeedsCleaning = true;
this._isBindingNeeded = true;
this._indexingNeeded = true;
this._moduleSymbolTable = undefined;
this._binderResults = undefined;
this._cachedIndexResults = undefined;
// If the file contains a wildcard import or __all__ symbols,
// we need to rebind because a dependent import may have changed.
if (this._parseResults) {
if (
this._parseResults.containsWildcardImport ||
AnalyzerNodeInfo.getDunderAllNames(this._parseResults.parseTree)
) {
this._parseTreeNeedsCleaning = true;
this._isBindingNeeded = true;
this._indexingNeeded = true;
this._moduleSymbolTable = undefined;
this._binderResults = undefined;
this._cachedIndexResults = undefined;
}
}
}

View File

@ -40,50 +40,3 @@ export function isTypedDictMemberAccessedThroughIndex(symbol: Symbol): boolean {
export function isFinalVariable(symbol: Symbol): boolean {
return symbol.getDeclarations().some((decl) => isFinalVariableDeclaration(decl));
}
export function getNamesInDunderAll(symbolTable: SymbolTable): string[] | undefined {
const namesToImport: string[] = [];
const allSymbol = 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
) {
namesToImport.push(listEntryNode.strings[0].value);
}
});
return namesToImport;
} else if (expr.nodeType === ParseNodeType.Tuple) {
expr.expressions.forEach((tupleEntryNode) => {
if (
tupleEntryNode.nodeType === ParseNodeType.StringList &&
tupleEntryNode.strings.length === 1 &&
tupleEntryNode.strings[0].nodeType === ParseNodeType.String
) {
namesToImport.push(tupleEntryNode.strings[0].value);
}
});
return namesToImport;
}
}
}
}
return undefined;
}

View File

@ -16,7 +16,7 @@ import { ImportLookup } from '../analyzer/analyzerFileInfo';
import * as AnalyzerNodeInfo from '../analyzer/analyzerNodeInfo';
import { AliasDeclaration, Declaration, DeclarationType } from '../analyzer/declaration';
import { getNameFromDeclaration } from '../analyzer/declarationUtils';
import { getLastTypedDeclaredForSymbol, getNamesInDunderAll } from '../analyzer/symbolUtils';
import { getLastTypedDeclaredForSymbol } from '../analyzer/symbolUtils';
import { TypeEvaluator } from '../analyzer/typeEvaluator';
import { isProperty } from '../analyzer/typeUtils';
import { throwIfCancellationRequested } from '../common/cancellationUtils';
@ -263,7 +263,7 @@ function collectSymbolIndexData(
const file = AnalyzerNodeInfo.getFileInfo(parseResults.parseTree);
let allNameTable: Set<string> | undefined;
if (autoImportMode && !file?.isStubFile) {
allNameTable = new Set<string>(getNamesInDunderAll(scope.symbolTable) ?? []);
allNameTable = new Set<string>(AnalyzerNodeInfo.getDunderAllNames(parseResults.parseTree) ?? []);
}
const symbolTable = scope.symbolTable;

View File

@ -0,0 +1,56 @@
/// <reference path="fourslash.ts" />
// @filename: testpkg/py.typed
// @library: true
////
// @filename: testpkg/__init__.py
// @library: true
//// from . import submod
//// from submod import foofoofoo5, foofoofoo6, foofoofoo7, foofoofoo8
//// foofoofoo1: int = 1
//// foofoofoo2: int = 2
//// foofoofoo3: int = 3
//// foofoofoo4: int = 4
//// __all__ = ["foofoofoo1"]
//// __all__ += ["foofoofoo2"]
//// __all__.extend(["foofoofoo3"])
//// __all__.extend(submod.__all__)
//// __all__.remove("foofoofoo1")
//// __all__.remove("foofoofoo6")
// @filename: testpkg/submod.py
// @library: true
//// foofoofoo5: int = 5
//// foofoofoo6: int = 6
//// foofoofoo7: int = 7
//// foofoofoo8: int = 8
//// __all__ = ["foofoofoo5"]
//// __all__ += ["foofoofoo6"]
//// __all__.extend(["foofoofoo7"])
// @filename: .src/test.py
//// from testpkg import *
//// foofoofoo[|/*marker1*/|]
// Ensure that only the __all__ items appear in the list.
// @ts-ignore
await helper.verifyCompletion('exact', {
marker1: {
completions: [
{
label: 'foofoofoo2',
},
{
label: 'foofoofoo3',
},
{
label: 'foofoofoo5',
},
{
label: 'foofoofoo7',
},
],
},
});