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

Added support for "eq" and "order" parameters in dataclass decorator as defined in PEP 557.

Browse Source
This commit is contained in:
Eric Traut 2020-11-19 15:25:26 -08:00
parent 37f926f6f0
commit 05e118a4b3
3 changed files with 86 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -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;
}
}

42
packages/pyright-internal/src/analyzer/types.ts
View File

@ -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) {

33
packages/pyright-internal/src/tests/samples/dataclass5.py
View File

@ -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)