mirror of
https://github.com/microsoft/pyright.git
synced 2024-09-11 16:06:39 +03:00
Added support for "eq" and "order" parameters in dataclass decorator as defined in PEP 557.
This commit is contained in:
parent
37f926f6f0
commit
05e118a4b3
@ -1994,22 +1994,30 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
|
||||
symbolTable.set('__new__', Symbol.createWithType(SymbolFlags.ClassMember, newType));
|
||||
}
|
||||
|
||||
// Synthesize comparison operators.
|
||||
['__eq__', '__ne__', '__lt__', '__le__', '__gt__', '__ge__'].forEach((operator) => {
|
||||
const synthesizeComparisonMethod = (operator: string, paramType: Type) => {
|
||||
const operatorMethod = FunctionType.createInstance(operator, '', FunctionTypeFlags.SynthesizedMethod);
|
||||
FunctionType.addParameter(operatorMethod, selfParam);
|
||||
FunctionType.addParameter(operatorMethod, {
|
||||
category: ParameterCategory.Simple,
|
||||
name: 'x',
|
||||
type:
|
||||
operator === '__eq__' || operator === '__ne__'
|
||||
? getBuiltInObject(node, 'object')
|
||||
: ObjectType.create(classType),
|
||||
type: paramType,
|
||||
hasDeclaredType: true,
|
||||
});
|
||||
operatorMethod.details.declaredReturnType = getBuiltInObject(node, 'bool');
|
||||
symbolTable.set(operator, Symbol.createWithType(SymbolFlags.ClassMember, operatorMethod));
|
||||
});
|
||||
};
|
||||
|
||||
// Synthesize comparison operators.
|
||||
if (!ClassType.isSkipSynthesizedDataclassEq(classType)) {
|
||||
synthesizeComparisonMethod('__eq__', getBuiltInObject(node, 'object'));
|
||||
}
|
||||
|
||||
if (!ClassType.isSkipSynthesizedDataclassOrder(classType)) {
|
||||
const objType = ObjectType.create(classType);
|
||||
['__lt__', '__le__', '__gt__', '__ge__'].forEach((operator) => {
|
||||
synthesizeComparisonMethod(operator, objType);
|
||||
});
|
||||
}
|
||||
|
||||
let dictType = getBuiltInType(node, 'Dict');
|
||||
if (isClass(dictType)) {
|
||||
@ -10014,7 +10022,7 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
|
||||
}
|
||||
|
||||
if (ClassType.isDataClass(classType)) {
|
||||
let skipSynthesizedInit = ClassType.isSkipSynthesizedInit(classType);
|
||||
let skipSynthesizedInit = ClassType.isSkipSynthesizedDataclassInit(classType);
|
||||
if (!skipSynthesizedInit) {
|
||||
// See if there's already a non-synthesized __init__ method.
|
||||
// We shouldn't override it.
|
||||
@ -10071,12 +10079,9 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
|
||||
decoratorCallType.category === TypeCategory.OverloadedFunction &&
|
||||
decoratorCallType.overloads[0].details.builtInName === 'dataclass'
|
||||
) {
|
||||
// Determine whether we should skip synthesizing the init method.
|
||||
let skipSynthesizeInit = false;
|
||||
|
||||
if (decoratorNode.expression.arguments) {
|
||||
decoratorNode.expression.arguments.forEach((arg) => {
|
||||
if (arg.name && arg.name.value === 'init') {
|
||||
if (arg.name) {
|
||||
if (arg.valueExpression) {
|
||||
const fileInfo = getFileInfo(decoratorNode);
|
||||
const value = evaluateStaticBoolExpression(
|
||||
@ -10084,7 +10089,16 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
|
||||
fileInfo.executionEnvironment
|
||||
);
|
||||
if (!value) {
|
||||
skipSynthesizeInit = true;
|
||||
if (arg.name.value === 'init') {
|
||||
originalClassType.details.flags |=
|
||||
ClassTypeFlags.SkipSynthesizedDataclassInit;
|
||||
} else if (arg.name.value === 'eq') {
|
||||
originalClassType.details.flags |=
|
||||
ClassTypeFlags.SkipSynthesizedDataclassEq;
|
||||
} else if (arg.name.value === 'order') {
|
||||
originalClassType.details.flags |=
|
||||
ClassTypeFlags.SkipSynthesizedDataclassOrder;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -10092,9 +10106,6 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
|
||||
}
|
||||
|
||||
originalClassType.details.flags |= ClassTypeFlags.DataClass;
|
||||
if (skipSynthesizeInit) {
|
||||
originalClassType.details.flags |= ClassTypeFlags.SkipSynthesizedInit;
|
||||
}
|
||||
return inputClassType;
|
||||
}
|
||||
}
|
||||
|
@ -239,65 +239,67 @@ export const enum ClassTypeFlags {
|
||||
|
||||
// Flags that control whether methods should be
|
||||
// synthesized for a dataclass class.
|
||||
SkipSynthesizedInit = 1 << 3,
|
||||
SkipSynthesizedDataclassInit = 1 << 3,
|
||||
SkipSynthesizedDataclassEq = 1 << 4,
|
||||
SkipSynthesizedDataclassOrder = 1 << 5,
|
||||
|
||||
// Introduced in PEP 589, TypedDict classes provide a way
|
||||
// to specify type hints for dictionaries with different
|
||||
// value types and a limited set of static keys.
|
||||
TypedDictClass = 1 << 4,
|
||||
TypedDictClass = 1 << 6,
|
||||
|
||||
// Used in conjunction with TypedDictClass, indicates that
|
||||
// the dictionary values can be omitted.
|
||||
CanOmitDictValues = 1 << 5,
|
||||
CanOmitDictValues = 1 << 7,
|
||||
|
||||
// The class has a metaclass of EnumMet or derives from
|
||||
// a class that has this metaclass.
|
||||
EnumClass = 1 << 6,
|
||||
EnumClass = 1 << 8,
|
||||
|
||||
// The class derives from a class that has the ABCMeta
|
||||
// metaclass. Such classes are allowed to contain
|
||||
// @abstractmethod decorators.
|
||||
SupportsAbstractMethods = 1 << 7,
|
||||
SupportsAbstractMethods = 1 << 9,
|
||||
|
||||
// The class has at least one abstract method or derives
|
||||
// from a base class that is abstract without providing
|
||||
// non-abstract overrides for all abstract methods.
|
||||
HasAbstractMethods = 1 << 8,
|
||||
HasAbstractMethods = 1 << 10,
|
||||
|
||||
// Derives from property class and has the semantics of
|
||||
// a property (with optional setter, deleter).
|
||||
PropertyClass = 1 << 9,
|
||||
PropertyClass = 1 << 11,
|
||||
|
||||
// The class is decorated with a "@final" decorator
|
||||
// indicating that it cannot be subclassed.
|
||||
Final = 1 << 10,
|
||||
Final = 1 << 12,
|
||||
|
||||
// The class derives directly from "Protocol".
|
||||
ProtocolClass = 1 << 11,
|
||||
ProtocolClass = 1 << 13,
|
||||
|
||||
// A class whose constructor (__init__ method) does not have
|
||||
// annotated types and is treated as though each parameter
|
||||
// is a generic type for purposes of type inference.
|
||||
PseudoGenericClass = 1 << 12,
|
||||
PseudoGenericClass = 1 << 14,
|
||||
|
||||
// A protocol class that is "runtime checkable" can be used
|
||||
// in an isinstance call.
|
||||
RuntimeCheckable = 1 << 13,
|
||||
RuntimeCheckable = 1 << 15,
|
||||
|
||||
// The type is defined in the typing_extensions.pyi file.
|
||||
TypingExtensionClass = 1 << 14,
|
||||
TypingExtensionClass = 1 << 16,
|
||||
|
||||
// The class type is in the process of being constructed and
|
||||
// is not yet complete. This allows us to detect cases where
|
||||
// the class refers to itself (e.g. uses itself as a type
|
||||
// argument to one of its generic base classes).
|
||||
PartiallyConstructed = 1 << 15,
|
||||
PartiallyConstructed = 1 << 17,
|
||||
|
||||
// The class or one of its ancestors defines a __class_getitem__
|
||||
// method that is used for subscripting. This is not set if the
|
||||
// class is generic, and therefore supports standard subscripting
|
||||
// semantics.
|
||||
HasCustomClassGetItem = 1 << 16,
|
||||
HasCustomClassGetItem = 1 << 18,
|
||||
}
|
||||
|
||||
interface ClassDetails {
|
||||
@ -507,8 +509,16 @@ export namespace ClassType {
|
||||
return !!(classType.details.flags & ClassTypeFlags.DataClass);
|
||||
}
|
||||
|
||||
export function isSkipSynthesizedInit(classType: ClassType) {
|
||||
return !!(classType.details.flags & ClassTypeFlags.SkipSynthesizedInit);
|
||||
export function isSkipSynthesizedDataclassInit(classType: ClassType) {
|
||||
return !!(classType.details.flags & ClassTypeFlags.SkipSynthesizedDataclassInit);
|
||||
}
|
||||
|
||||
export function isSkipSynthesizedDataclassEq(classType: ClassType) {
|
||||
return !!(classType.details.flags & ClassTypeFlags.SkipSynthesizedDataclassEq);
|
||||
}
|
||||
|
||||
export function isSkipSynthesizedDataclassOrder(classType: ClassType) {
|
||||
return !!(classType.details.flags & ClassTypeFlags.SkipSynthesizedDataclassOrder);
|
||||
}
|
||||
|
||||
export function isTypedDictClass(classType: ClassType) {
|
||||
|
@ -2,6 +2,8 @@
|
||||
# with a custom __init__.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal
|
||||
|
||||
|
||||
@dataclass(init=False)
|
||||
class A:
|
||||
@ -12,8 +14,10 @@ class A:
|
||||
self.x = x
|
||||
self.x_squared = x ** 2
|
||||
|
||||
|
||||
a = A(3)
|
||||
|
||||
|
||||
@dataclass(init=True)
|
||||
class B:
|
||||
x: int
|
||||
@ -23,8 +27,10 @@ class B:
|
||||
self.x = x
|
||||
self.x_squared = x ** 2
|
||||
|
||||
|
||||
b = B(3)
|
||||
|
||||
|
||||
@dataclass()
|
||||
class C:
|
||||
x: int
|
||||
@ -34,13 +40,40 @@ class C:
|
||||
self.x = x
|
||||
self.x_squared = x ** 2
|
||||
|
||||
|
||||
c = C(3)
|
||||
|
||||
|
||||
@dataclass(init=False)
|
||||
class D:
|
||||
x: int
|
||||
x_squared: int
|
||||
|
||||
|
||||
# This should generate an error because there is no
|
||||
# override __init__ method and no synthesized __init__.
|
||||
d = D(3)
|
||||
|
||||
|
||||
@dataclass(eq=False)
|
||||
class E:
|
||||
x: int
|
||||
|
||||
def __eq__(self, x: "E") -> float:
|
||||
return 1.23
|
||||
|
||||
|
||||
foo1 = E(3) == E(3)
|
||||
t1: Literal["float"] = reveal_type(foo1)
|
||||
|
||||
|
||||
@dataclass(order=False)
|
||||
class F:
|
||||
x: int
|
||||
|
||||
def __lt__(self, x: "F") -> float:
|
||||
return 1.23
|
||||
|
||||
|
||||
foo2 = F(3) < F(3)
|
||||
t1: Literal["float"] = reveal_type(foo2)
|
||||
|
Loading…
Reference in New Issue
Block a user