microsoft/pyright
1
1
Fork 0
mirror of https://github.com/microsoft/pyright.git synced 2024-09-11 07:55:56 +03:00
Code Issues Projects Releases Wiki Activity

Fixed bug that results in incorrect type evaluation behavior within c… (#7912)

Browse Source 
* Fixed bug that results in incorrect type evaluation behavior within class body of an enum class when one enum member is used to define another enum member. This involved a significant rewrite of the logic involving enum symbol evaluation. It addresses #7763.

* Add recursion protection.
This commit is contained in:
Eric Traut 2024-05-13 21:58:47 -07:00 committed by GitHub
parent c477556fa1
commit 506f2c14dd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 160 additions and 143 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
packages/pyright-internal/src/analyzer/checker.ts
View File

@ -98,7 +98,7 @@ import { getBoundCallMethod, getBoundInitMethod, getBoundNewMethod } from './con
import { Declaration, DeclarationType, isAliasDeclaration } from './declaration';
import { getNameNodeForDeclaration } from './declarationUtils';
import { deprecatedAliases, deprecatedSpecialForms } from './deprecatedSymbols';
import { getEnumDeclaredValueType, isEnumClassWithMembers } from './enums';
import { getEnumDeclaredValueType, isEnumClassWithMembers, transformTypeForEnumMember } from './enums';
import { ImportResolver, ImportedModuleDescriptor, createImportedModuleDescriptor } from './importResolver';
import { ImportResult, ImportType } from './importResult';
import { getRelativeModuleName, getTopLevelImports } from './importStatementUtils';
@ -4977,9 +4977,11 @@ export class Checker extends ParseTreeWalker {
return;
}
const symbolType = this._evaluator.getEffectiveTypeOfSymbol(symbol);
const symbolType = transformTypeForEnumMember(this._evaluator, classType, name);
// Is this symbol a literal instance of the enum class?
if (
!symbolType ||
!isClassInstance(symbolType) ||
!ClassType.isSameGenericClass(symbolType, classType) ||
!(symbolType.literalValue instanceof EnumLiteral)

153
packages/pyright-internal/src/analyzer/enums.ts
View File

@ -11,16 +11,10 @@ import { assert } from '../common/debug';
import { ArgumentCategory, ExpressionNode, NameNode, ParseNode, ParseNodeType } from '../parser/parseNodes';
import { getFileInfo } from './analyzerNodeInfo';
import { VariableDeclaration } from './declaration';
import {
getClassFullName,
getEnclosingClass,
getParentNodeOfType,
getTypeSourceId,
isNodeContainedWithin,
} from './parseTreeUtils';
import { getClassFullName, getEnclosingClass, getTypeSourceId } from './parseTreeUtils';
import { Symbol, SymbolFlags } from './symbol';
import { isPrivateName, isSingleDunderName } from './symbolNameUtils';
import { FunctionArgument, TypeEvaluator, TypeResult } from './typeEvaluatorTypes';
import { EvaluatorFlags, FunctionArgument, TypeEvaluator, TypeResult } from './typeEvaluatorTypes';
import { enumerateLiteralsForType } from './typeGuards';
import { MemberAccessFlags, computeMroLinearization, lookUpClassMember, makeInferenceContext } from './typeUtils';
import {
@ -29,6 +23,7 @@ import {
ClassTypeFlags,
EnumLiteral,
Type,
TypeBase,
UnknownType,
combineTypes,
findSubtype,
@ -38,6 +33,7 @@ import {
isFunction,
isInstantiableClass,
isOverloadedFunction,
maxTypeRecursionCount,
} from './types';
// Determines whether the class is an Enum metaclass or a subclass thereof.
@ -55,16 +51,16 @@ export function isEnumClassWithMembers(evaluator: TypeEvaluator, classType: Clas
}
// Determine whether the enum class defines a member.
let definesValue = false;
let definesMember = false;
ClassType.getSymbolTable(classType).forEach((symbol) => {
const symbolType = evaluator.getEffectiveTypeOfSymbol(symbol);
if (isClassInstance(symbolType) && ClassType.isSameGenericClass(symbolType, classType)) {
definesValue = true;
ClassType.getSymbolTable(classType).forEach((symbol, name) => {
const symbolType = transformTypeForEnumMember(evaluator, classType, name);
if (symbolType && isClassInstance(symbolType) && ClassType.isSameGenericClass(symbolType, classType)) {
definesMember = true;
}
});
return definesValue;
return definesMember;
}
// Creates a new custom enum class with named values.
@ -287,24 +283,33 @@ export function createEnumType(
return classType;
}
export function transformTypeForPossibleEnumClass(
export function transformTypeForEnumMember(
evaluator: TypeEvaluator,
statementNode: ParseNode,
nameNode: NameNode,
getValueType: () => { declaredType?: Type; assignedType?: Type }
classType: ClassType,
memberName: string,
recursionCount = 0
): Type | undefined {
// If the node is within a class that derives from the metaclass
// "EnumMeta", we need to treat assignments differently.
const enclosingClassNode = getEnclosingClass(statementNode, /* stopAtFunction */ true);
if (!enclosingClassNode) {
if (recursionCount > maxTypeRecursionCount) {
return undefined;
}
recursionCount++;
if (!ClassType.isEnumClass(classType)) {
return undefined;
}
const enumClassInfo = evaluator.getTypeOfClass(enclosingClassNode);
if (!enumClassInfo || !ClassType.isEnumClass(enumClassInfo.classType)) {
const memberInfo = lookUpClassMember(classType, memberName);
if (!memberInfo || !isClass(memberInfo.classType) || !ClassType.isEnumClass(memberInfo.classType)) {
return undefined;
}
const decls = memberInfo.symbol.getDeclarations();
if (decls.length < 1) {
return undefined;
}
const primaryDecl = decls[0];
// In ".py" files, the transform applies only to members that are
// assigned within the class. In stub files, it applies to most variables
// even if they are not assigned. This unfortunate convention means
@ -314,22 +319,36 @@ export function transformTypeForPossibleEnumClass(
// it, we're stuck with this limitation.
let isMemberOfEnumeration = false;
let isUnpackedTuple = false;
let valueTypeExprNode: ExpressionNode | undefined;
let declaredTypeNode: ExpressionNode | undefined;
let nameNode: NameNode | undefined;
if (
statementNode.nodeType === ParseNodeType.Assignment &&
isNodeContainedWithin(nameNode, statementNode.leftExpression)
if (primaryDecl.node.nodeType === ParseNodeType.Name) {
nameNode = primaryDecl.node;
} else if (primaryDecl.node.nodeType === ParseNodeType.Function) {
// Handle the case where a method is decorated with @enum.member.
nameNode = primaryDecl.node.name;
} else {
return undefined;
}
if (nameNode.parent?.nodeType === ParseNodeType.Assignment && nameNode.parent.leftExpression === nameNode) {
isMemberOfEnumeration = true;
valueTypeExprNode = nameNode.parent.rightExpression;
} else if (
nameNode.parent?.nodeType === ParseNodeType.Tuple &&
nameNode.parent.parent?.nodeType === ParseNodeType.Assignment
) {
isMemberOfEnumeration = true;
if (getParentNodeOfType(nameNode, ParseNodeType.Tuple)) {
isUnpackedTuple = true;
}
isUnpackedTuple = true;
valueTypeExprNode = nameNode.parent.parent.rightExpression;
} else if (
getFileInfo(nameNode).isStubFile &&
nameNode.parent?.nodeType === ParseNodeType.TypeAnnotation &&
nameNode.parent.valueExpression === nameNode
) {
isMemberOfEnumeration = true;
declaredTypeNode = nameNode.parent.typeAnnotation;
}
// The spec specifically excludes names that start and end with a single underscore.
@ -343,9 +362,39 @@ export function transformTypeForPossibleEnumClass(
return undefined;
}
const valueTypeInfo = getValueType();
const declaredType = valueTypeInfo.declaredType;
let assignedType = valueTypeInfo.assignedType;
const declaredType = declaredTypeNode ? evaluator.getTypeOfAnnotation(declaredTypeNode) : undefined;
let assignedType: Type | undefined;
if (valueTypeExprNode) {
const evalFlags = getFileInfo(valueTypeExprNode).isStubFile ? EvaluatorFlags.ConvertEllipsisToAny : undefined;
assignedType = evaluator.getTypeOfExpression(valueTypeExprNode, evalFlags).type;
}
// Handle aliases to other enum members within the same enum.
if (valueTypeExprNode?.nodeType === ParseNodeType.Name && valueTypeExprNode.value !== memberName) {
const aliasedEnumType = transformTypeForEnumMember(
evaluator,
classType,
valueTypeExprNode.value,
recursionCount
);
if (
aliasedEnumType &&
isClassInstance(aliasedEnumType) &&
ClassType.isSameGenericClass(aliasedEnumType, ClassType.cloneAsInstance(memberInfo.classType)) &&
aliasedEnumType.literalValue !== undefined
) {
return aliasedEnumType;
}
}
if (primaryDecl.node.nodeType === ParseNodeType.Function) {
const functionType = evaluator.getTypeOfFunction(primaryDecl.node);
if (functionType) {
assignedType = functionType.decoratedType;
}
}
let valueType = declaredType ?? assignedType ?? UnknownType.create();
@ -377,9 +426,13 @@ export function transformTypeForPossibleEnumClass(
return undefined;
}
if (!assignedType && statementNode.nodeType === ParseNodeType.Assignment) {
if (
!assignedType &&
nameNode.parent?.nodeType === ParseNodeType.Assignment &&
nameNode.parent.leftExpression === nameNode
) {
assignedType = evaluator.getTypeOfExpression(
statementNode.rightExpression,
nameNode.parent.rightExpression,
/* flags */ undefined,
makeInferenceContext(declaredType)
).type;
@ -411,27 +464,18 @@ export function transformTypeForPossibleEnumClass(
}
}
// Handle aliases to other enum members within the same enum.
if (
isClassInstance(valueType) &&
ClassType.isSameGenericClass(valueType, ClassType.cloneAsInstance(enumClassInfo.classType)) &&
valueType.literalValue !== undefined
) {
return undefined;
}
if (!isMemberOfEnumeration) {
return undefined;
}
const enumLiteral = new EnumLiteral(
enumClassInfo.classType.details.fullName,
enumClassInfo.classType.details.name,
memberInfo.classType.details.fullName,
memberInfo.classType.details.name,
nameNode.value,
valueType
);
return ClassType.cloneAsInstance(ClassType.cloneWithLiteral(enumClassInfo.classType, enumLiteral));
return ClassType.cloneAsInstance(ClassType.cloneWithLiteral(memberInfo.classType, enumLiteral));
}
export function isDeclInEnumClass(evaluator: TypeEvaluator, decl: VariableDeclaration): boolean {
@ -483,11 +527,20 @@ export function getTypeOfEnumMember(
memberName: string,
isIncomplete: boolean
): TypeResult | undefined {
// Handle the special case of 'name' and 'value' members within an enum.
if (!isClassInstance(classType) || !ClassType.isEnumClass(classType)) {
if (!ClassType.isEnumClass(classType)) {
return undefined;
}
if (TypeBase.isInstantiable(classType)) {
const type = transformTypeForEnumMember(evaluator, classType, memberName);
if (!type) {
return undefined;
}
return { type, isIncomplete };
}
// Handle the special case of 'name' and 'value' members within an enum.
const literalValue = classType.literalValue;
if (memberName === 'name' || memberName === '_name_') {

103
packages/pyright-internal/src/analyzer/typeEvaluator.ts
View File

@ -138,7 +138,6 @@ import {
isDeclInEnumClass,
isEnumClassWithMembers,
isEnumMetaclass,
transformTypeForPossibleEnumClass,
} from './enums';
import { applyFunctionTransform } from './functionTransform';
import { createNamedTupleType } from './namedTuples';
@ -5329,8 +5328,35 @@ export function createTypeEvaluator(
case TypeCategory.Class: {
let typeResult: TypeResult | undefined;
if (usage.method === 'get') {
typeResult = getTypeOfEnumMember(evaluatorInterface, node, baseType, memberName, isIncomplete);
const enumMemberResult = getTypeOfEnumMember(
evaluatorInterface,
node,
baseType,
memberName,
isIncomplete
);
if (enumMemberResult) {
if (usage.method === 'get') {
typeResult = enumMemberResult;
} else {
// Is this an attempt to delete or overwrite an enum member?
if (
isClassInstance(enumMemberResult.type) &&
ClassType.isSameGenericClass(enumMemberResult.type, baseType) &&
enumMemberResult.type.literalValue !== undefined
) {
const diagMessage =
usage.method === 'set' ? LocMessage.enumMemberSet() : LocMessage.enumMemberDelete();
addDiagnostic(
DiagnosticRule.reportAttributeAccessIssue,
diagMessage.format({ name: memberName }) + diag.getString(),
node.memberName,
diag.getEffectiveTextRange() ?? node.memberName
);
}
}
}
if (!typeResult) {
@ -15985,23 +16011,6 @@ export function createTypeEvaluator(
// If this is an enum, transform the type as required.
rightHandType = srcType;
let targetName: NameNode | undefined;
if (node.leftExpression.nodeType === ParseNodeType.Name) {
targetName = node.leftExpression;
} else if (
node.leftExpression.nodeType === ParseNodeType.TypeAnnotation &&
node.leftExpression.valueExpression.nodeType === ParseNodeType.Name
) {
targetName = node.leftExpression.valueExpression;
}
if (targetName) {
rightHandType =
transformTypeForPossibleEnumClass(evaluatorInterface, node, targetName, () => {
return { assignedType: rightHandType };
}) ?? rightHandType;
}
if (typeAliasNameNode) {
// If this was a speculative type alias, it becomes a real type alias
// only if the evaluated type is an instantiable type.
@ -17558,12 +17567,6 @@ export function createTypeEvaluator(
}
}
// In case this is an enum class and a method wrapped in an enum.member.
decoratedType =
transformTypeForPossibleEnumClass(evaluatorInterface, node, node.name, () => {
return { assignedType: decoratedType };
}) ?? decoratedType;
// See if there are any overloads provided by previous function declarations.
if (isFunction(decoratedType)) {
decoratedType.details.deprecatedMessage = functionType.details.deprecatedMessage;
@ -20973,25 +20976,6 @@ export function createTypeEvaluator(
}
if (declaredType) {
// Apply enum transform if appropriate.
if (declaration.node.nodeType === ParseNodeType.Name) {
const variableNode =
ParseTreeUtils.getParentNodeOfType(declaration.node, ParseNodeType.Assignment) ??
ParseTreeUtils.getParentNodeOfType(declaration.node, ParseNodeType.TypeAnnotation);
if (variableNode) {
declaredType =
transformTypeForPossibleEnumClass(
evaluatorInterface,
variableNode,
declaration.node,
() => {
return { declaredType };
}
) ?? declaredType;
}
}
if (isClassInstance(declaredType) && ClassType.isBuiltIn(declaredType, 'TypeAlias')) {
return { type: undefined, isTypeAlias: true };
}
@ -21311,35 +21295,6 @@ export function createTypeEvaluator(
evaluateTypesForStatement(typeSource);
})?.type;
if (inferredType && resolvedDecl.node.nodeType === ParseNodeType.Name) {
const variableNode =
ParseTreeUtils.getParentNodeOfType(resolvedDecl.node, ParseNodeType.Assignment) ??
ParseTreeUtils.getParentNodeOfType(resolvedDecl.node, ParseNodeType.TypeAnnotation);
if (variableNode) {
// See if this is an enum member. If so, we need to handle it as a special case.
const enumMemberType = transformTypeForPossibleEnumClass(
evaluatorInterface,
variableNode,
resolvedDecl.node,
() => {
assert(resolvedDecl.inferredTypeSource !== undefined);
const inferredTypeSource = resolvedDecl.inferredTypeSource;
return {
assignedType:
evaluateTypeForSubnode(inferredTypeSource, () => {
evaluateTypesForStatement(inferredTypeSource);
})?.type ?? UnknownType.create(),
};
}
);
if (enumMemberType) {
inferredType = enumMemberType;
}
}
}
if (inferredType && isTypeAlias && resolvedDecl.typeAliasName) {
// If this was a speculative type alias, it becomes a real type alias only
// in the event that its inferred type is instantiable or explicitly Any

7
packages/pyright-internal/src/analyzer/typeGuards.ts
View File

@ -24,6 +24,7 @@ import { KeywordType, OperatorType } from '../parser/tokenizerTypes';
import { getFileInfo } from './analyzerNodeInfo';
import { populateTypeVarContextBasedOnExpectedType } from './constraintSolver';
import { Declaration, DeclarationType } from './declaration';
import { transformTypeForEnumMember } from './enums';
import * as ParseTreeUtils from './parseTreeUtils';
import { ScopeType } from './scope';
import { getScopeForNode } from './scopeUtils';
@ -2574,9 +2575,11 @@ export function enumerateLiteralsForType(evaluator: TypeEvaluator, type: ClassTy
// Enumerate all of the values in this enumeration.
const enumList: ClassType[] = [];
const fields = ClassType.getSymbolTable(type);
fields.forEach((symbol) => {
fields.forEach((symbol, name) => {
if (!symbol.isIgnoredForProtocolMatch()) {
const symbolType = evaluator.getEffectiveTypeOfSymbol(symbol);
let symbolType = evaluator.getEffectiveTypeOfSymbol(symbol);
symbolType = transformTypeForEnumMember(evaluator, type, name) ?? symbolType;
if (
isClassInstance(symbolType) &&
ClassType.isSameGenericClass(type, symbolType) &&

4
packages/pyright-internal/src/localization/localize.ts
View File

@ -432,6 +432,10 @@ export namespace Localizer {
export const ellipsisSecondArg = () => getRawString('Diagnostic.ellipsisSecondArg');
export const enumClassOverride = () =>
new ParameterizedString<{ name: string }>(getRawString('Diagnostic.enumClassOverride'));
export const enumMemberDelete = () =>
new ParameterizedString<{ name: string }>(getRawString('Diagnostic.enumMemberDelete'));
export const enumMemberSet = () =>
new ParameterizedString<{ name: string }>(getRawString('Diagnostic.enumMemberSet'));
export const exceptionGroupIncompatible = () => getRawString('Diagnostic.exceptionGroupIncompatible');
export const exceptionTypeIncorrect = () =>
new ParameterizedString<{ type: string }>(getRawString('Diagnostic.exceptionTypeIncorrect'));

2
packages/pyright-internal/src/localization/package.nls.en-us.json
View File

@ -140,6 +140,8 @@
"ellipsisContext": "\"...\" is not allowed in this context",
"ellipsisSecondArg": "\"...\" is allowed only as the second of two arguments",
"enumClassOverride": "Enum class \"{name}\" is final and cannot be subclassed",
"enumMemberDelete": "Enum member \"{name}\" cannot be deleted",
"enumMemberSet": "Enum member \"{name}\" cannot be assigned",
"exceptionGroupIncompatible": "Exception group syntax (\"except*\") requires Python 3.11 or newer",
"exceptionTypeIncorrect": "\"{type}\" does not derive from BaseException",
"exceptionTypeNotClass": "\"{type}\" is not a valid exception class",

4
packages/pyright-internal/src/tests/completions.test.ts
View File

@ -8,9 +8,9 @@ import assert from 'assert';
import { CancellationToken } from 'vscode-languageserver';
import { CompletionItemKind, MarkupKind } from 'vscode-languageserver-types';
import { Uri } from '../common/uri/uri';
import { CompletionOptions, CompletionProvider } from '../languageService/completionProvider';
import { parseAndGetTestState } from './harness/fourslash/testState';
import { Uri } from '../common/uri/uri';
test('completion import statement tooltip', async () => {
const code = `
@ -1148,7 +1148,7 @@ test('Enum member', async () => {
{
label: 'this',
kind: CompletionItemKind.EnumMember,
documentation: '```python\nthis: Literal[MyEnum.this]\n```',
documentation: '```python\nthis: int\n```',
},
],
},

11
packages/pyright-internal/src/tests/samples/enum1.py
View File

@ -249,3 +249,14 @@ class TestEnum19(Enum):
reveal_type(TestEnum19.A, expected_text="Literal[TestEnum19.A]")
reveal_type(TestEnum19.__B, expected_text="Literal[2]")
class TestEnum20(Enum):
A = 1
B = A + 1
reveal_type(TestEnum20.A, expected_text="Literal[TestEnum20.A]")
reveal_type(TestEnum20.A.value, expected_text="Literal[1]")
reveal_type(TestEnum20.B, expected_text="Literal[TestEnum20.B]")
reveal_type(TestEnum20.B.value, expected_text="Literal[2]")

13
packages/pyright-internal/src/tests/samples/enum9.py
View File

@ -25,16 +25,3 @@ reveal_type(Enum1.NON_MEMBER, expected_text="int")
reveal_type(Enum1.MEMBER.value, expected_text="Literal[1]")
reveal_type(Enum1.ANOTHER_MEMBER.value, expected_text="int")
reveal_type(Enum1.ALSO_A_MEMBER.value, expected_text="() -> Literal[4]")
class Enum2(Enum):
MEMBER: int = member(1)
NON_MEMBER: int = nonmember(1)
NON_MEMBER2: ClassVar[dict[str, str | int]] = nonmember({})
reveal_type(Enum2.MEMBER, expected_text="Literal[Enum2.MEMBER]")
reveal_type(Enum2.NON_MEMBER, expected_text="int")
reveal_type(Enum2.NON_MEMBER2, expected_text="dict[str, str | int]")
reveal_type(Enum2.MEMBER.value, expected_text="int")