From 5ca499b1125497596dd4479f7bdfb56f2ebd00e5 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Mon, 28 Sep 2020 14:45:00 -0700 Subject: [PATCH] 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. --- .../src/analyzer/analyzerFileInfo.ts | 1 + .../src/analyzer/analyzerNodeInfo.ts | 13 ++ .../pyright-internal/src/analyzer/binder.ts | 179 ++++++++++++++++-- .../pyright-internal/src/analyzer/program.ts | 2 + .../src/analyzer/sourceFile.ts | 23 ++- .../src/analyzer/symbolUtils.ts | 47 ----- .../languageService/documentSymbolProvider.ts | 4 +- .../import.pytyped.dunderAll.fourslash.ts | 56 ++++++ 8 files changed, 248 insertions(+), 77 deletions(-) create mode 100644 packages/pyright-internal/src/tests/fourslash/import.pytyped.dunderAll.fourslash.ts diff --git a/packages/pyright-internal/src/analyzer/analyzerFileInfo.ts b/packages/pyright-internal/src/analyzer/analyzerFileInfo.ts index 4638b7249..450e72c8f 100644 --- a/packages/pyright-internal/src/analyzer/analyzerFileInfo.ts +++ b/packages/pyright-internal/src/analyzer/analyzerFileInfo.ts @@ -20,6 +20,7 @@ export type ImportLookup = (filePath: string) => ImportLookupResult | undefined; export interface ImportLookupResult { symbolTable: SymbolTable; + dunderAllNames: string[] | undefined; docString?: string; } diff --git a/packages/pyright-internal/src/analyzer/analyzerNodeInfo.ts b/packages/pyright-internal/src/analyzer/analyzerNodeInfo.ts index 83f5526b3..280c8e913 100644 --- a/packages/pyright-internal/src/analyzer/analyzerNodeInfo.ts +++ b/packages/pyright-internal/src/analyzer/analyzerNodeInfo.ts @@ -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; + + // 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 = new Map(); + // 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([])"? + 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(.__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( diff --git a/packages/pyright-internal/src/analyzer/program.ts b/packages/pyright-internal/src/analyzer/program.ts index 59f498557..980d6fec2 100644 --- a/packages/pyright-internal/src/analyzer/program.ts +++ b/packages/pyright-internal/src/analyzer/program.ts @@ -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, }; }; diff --git a/packages/pyright-internal/src/analyzer/sourceFile.ts b/packages/pyright-internal/src/analyzer/sourceFile.ts index 4e40bc74e..fe12157a8 100644 --- a/packages/pyright-internal/src/analyzer/sourceFile.ts +++ b/packages/pyright-internal/src/analyzer/sourceFile.ts @@ -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; + } } } diff --git a/packages/pyright-internal/src/analyzer/symbolUtils.ts b/packages/pyright-internal/src/analyzer/symbolUtils.ts index 2c8ea3b44..a71d8b5bb 100644 --- a/packages/pyright-internal/src/analyzer/symbolUtils.ts +++ b/packages/pyright-internal/src/analyzer/symbolUtils.ts @@ -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; -} diff --git a/packages/pyright-internal/src/languageService/documentSymbolProvider.ts b/packages/pyright-internal/src/languageService/documentSymbolProvider.ts index d39aa68ee..90ec30022 100644 --- a/packages/pyright-internal/src/languageService/documentSymbolProvider.ts +++ b/packages/pyright-internal/src/languageService/documentSymbolProvider.ts @@ -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 | undefined; if (autoImportMode && !file?.isStubFile) { - allNameTable = new Set(getNamesInDunderAll(scope.symbolTable) ?? []); + allNameTable = new Set(AnalyzerNodeInfo.getDunderAllNames(parseResults.parseTree) ?? []); } const symbolTable = scope.symbolTable; diff --git a/packages/pyright-internal/src/tests/fourslash/import.pytyped.dunderAll.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/import.pytyped.dunderAll.fourslash.ts new file mode 100644 index 000000000..80a613d34 --- /dev/null +++ b/packages/pyright-internal/src/tests/fourslash/import.pytyped.dunderAll.fourslash.ts @@ -0,0 +1,56 @@ +/// + +// @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', + }, + ], + }, +});