def visit_type_var_expr(self, node: TypeVarExpr) -> Node:
return TypeVarExpr(
node.name(),
node.fullname(),
self.types(node.values),
self.type(node.upper_bound),
variance=node.variance)
def add_construct_method(self, fields: List["PydanticModelField"]) -> None:
"""
Adds a fully typed `construct` classmethod to the class.
Similar to the fields-aware __init__ method, but always uses the field names (not aliases),
and does not treat settings fields as optional.
"""
ctx = self._ctx
set_str = ctx.api.named_type("__builtins__.set",
[ctx.api.named_type("__builtins__.str")])
optional_set_str = UnionType([set_str, NoneType()])
fields_set_argument = Argument(Var("_fields_set", optional_set_str),
optional_set_str, None, ARG_OPT)
construct_arguments = self.get_field_arguments(
fields, typed=True, force_all_optional=False, use_alias=False)
construct_arguments = [fields_set_argument] + construct_arguments
obj_type = ctx.api.named_type("__builtins__.object")
self_tvar_name = "Model"
tvar_fullname = ctx.cls.fullname + "." + self_tvar_name
tvd = TypeVarDef(self_tvar_name, tvar_fullname, -1, [], obj_type)
self_tvar_expr = TypeVarExpr(self_tvar_name, tvar_fullname, [],
obj_type)
ctx.cls.info.names[self_tvar_name] = SymbolTableNode(
MDEF, self_tvar_expr)
self_type = TypeVarType(tvd)
add_method(
ctx,
"construct",
construct_arguments,
return_type=self_type,
self_type=self_type,
tvar_def=tvd,
is_classmethod=True,
)
def _add_cmp(ctx: 'mypy.plugin.ClassDefContext', adder: 'MethodAdder') -> None:
"""Generate all the cmp methods for this class."""
# For __ne__ and __eq__ the type is:
# def __ne__(self, other: object) -> bool
bool_type = ctx.api.named_type('__builtins__.bool')
object_type = ctx.api.named_type('__builtins__.object')
args = [Argument(Var('other', object_type), object_type, None, ARG_POS)]
for method in ['__ne__', '__eq__']:
adder.add_method(method, args, bool_type)
# For the rest we use:
# AT = TypeVar('AT')
# def __lt__(self: AT, other: AT) -> bool
# This way comparisons with subclasses will work correctly.
tvd = TypeVarDef(SELF_TVAR_NAME,
ctx.cls.info.fullname() + '.' + SELF_TVAR_NAME, -1, [],
object_type)
tvd_type = TypeVarType(tvd)
self_tvar_expr = TypeVarExpr(
SELF_TVAR_NAME,
ctx.cls.info.fullname() + '.' + SELF_TVAR_NAME, [], object_type)
ctx.cls.info.names[SELF_TVAR_NAME] = SymbolTableNode(MDEF, self_tvar_expr)
args = [Argument(Var('other', tvd_type), tvd_type, None, ARG_POS)]
for method in ['__lt__', '__le__', '__gt__', '__ge__']:
adder.add_method(method, args, bool_type, self_type=tvd_type, tvd=tvd)
def bind_new(self, name: str, tvar_expr: TypeVarExpr) -> TypeVarDef:
if self.is_class_scope:
self.class_id += 1
i = self.class_id
else:
self.func_id -= 1
i = self.func_id
tvar_def = TypeVarDef(name,
tvar_expr.fullname(),
i,
values=tvar_expr.values,
upper_bound=tvar_expr.upper_bound,
variance=tvar_expr.variance,
line=tvar_expr.line,
column=tvar_expr.column)
self.scope[tvar_expr.fullname()] = tvar_def
return tvar_def
def bind_new(self, name: str, tvar_expr: TypeVarExpr) -> TypeVarDef:
if self.is_class_scope:
self.class_id += 1
i = self.class_id
else:
self.func_id -= 1
i = self.func_id
tvar_def = TypeVarDef(name,
tvar_expr.fullname(),
i,
values=tvar_expr.values,
upper_bound=tvar_expr.upper_bound,
variance=tvar_expr.variance,
line=tvar_expr.line,
column=tvar_expr.column)
self.scope[tvar_expr.fullname()] = tvar_def
return tvar_def
def _add_typevar(context: ClassDefContext,
tVarName: str) -> mypy.types.TypeVarDef:
typeInfo = context.cls.info
tVarQualifiedName = f'{get_fullname(typeInfo)}.{tVarName}'
objectType = context.api.named_type('__builtins__.object')
tVarExpr = TypeVarExpr(tVarName, tVarQualifiedName, [], objectType)
typeInfo.names[tVarName] = SymbolTableNode(MDEF, tVarExpr)
return mypy.types.TypeVarDef(tVarName, tVarQualifiedName, -1, [],
objectType)
def _add_order(ctx: 'mypy.plugin.ClassDefContext', adder: 'MethodAdder') -> None:
"""Generate all the ordering methods for this class."""
bool_type = ctx.api.named_type('__builtins__.bool')
object_type = ctx.api.named_type('__builtins__.object')
# Make the types be:
# AT = TypeVar('AT')
# def __lt__(self: AT, other: AT) -> bool
# This way comparisons with subclasses will work correctly.
tvd = TypeVarType(SELF_TVAR_NAME, ctx.cls.info.fullname + '.' + SELF_TVAR_NAME,
-1, [], object_type)
self_tvar_expr = TypeVarExpr(SELF_TVAR_NAME, ctx.cls.info.fullname + '.' + SELF_TVAR_NAME,
[], object_type)
ctx.cls.info.names[SELF_TVAR_NAME] = SymbolTableNode(MDEF, self_tvar_expr)
args = [Argument(Var('other', tvd), tvd, None, ARG_POS)]
for method in ['__lt__', '__le__', '__gt__', '__ge__']:
adder.add_method(method, args, bool_type, self_type=tvd, tvd=tvd)
def add_construct_method(self, fields: List['PydanticModelField']) -> None:
"""
Adds a fully typed `construct` classmethod to the class.
Similar to the fields-aware __init__ method, but always uses the field names (not aliases),
and does not treat settings fields as optional.
"""
ctx = self._ctx
set_str = ctx.api.named_type(
f'{BUILTINS_NAME}.set',
[ctx.api.named_type(f'{BUILTINS_NAME}.str')])
optional_set_str = UnionType([set_str, NoneType()])
fields_set_argument = Argument(Var('_fields_set', optional_set_str),
optional_set_str, None, ARG_OPT)
construct_arguments = self.get_field_arguments(
fields, typed=True, force_all_optional=False, use_alias=False)
construct_arguments = [fields_set_argument] + construct_arguments
obj_type = ctx.api.named_type(f'{BUILTINS_NAME}.object')
self_tvar_name = '_PydanticBaseModel' # Make sure it does not conflict with other names in the class
tvar_fullname = ctx.cls.fullname + '.' + self_tvar_name
tvd = TypeVarDef(self_tvar_name, tvar_fullname, -1, [], obj_type)
self_tvar_expr = TypeVarExpr(self_tvar_name, tvar_fullname, [],
obj_type)
ctx.cls.info.names[self_tvar_name] = SymbolTableNode(
MDEF, self_tvar_expr)
# Backward-compatible with TypeVarDef from Mypy 0.910.
if isinstance(tvd, TypeVarType):
self_type = tvd
else:
self_type = TypeVarType(tvd) # type: ignore[call-arg]
add_method(
ctx,
'construct',
construct_arguments,
return_type=self_type,
self_type=self_type,
tvar_def=tvd,
is_classmethod=True,
)
def transform(self) -> None:
"""Apply all the necessary transformations to the underlying
dataclass so as to ensure it is fully type checked according
to the rules in PEP 557.
"""
ctx = self._ctx
info = self._ctx.cls.info
attributes = self.collect_attributes()
if attributes is None:
# Some definitions are not ready, defer() should be already called.
return
for attr in attributes:
if attr.type is None:
ctx.api.defer()
return
decorator_arguments = {
'init': _get_decorator_bool_argument(self._ctx, 'init', True),
'eq': _get_decorator_bool_argument(self._ctx, 'eq', True),
'order': _get_decorator_bool_argument(self._ctx, 'order', False),
'frozen': _get_decorator_bool_argument(self._ctx, 'frozen', False),
}
if info.get('replace') is None:
obj_type = ctx.api.named_type('__builtins__.object')
self_tvar_expr = TypeVarExpr(SELF_UVAR_NAME,
info.fullname + '.' + SELF_UVAR_NAME,
[], obj_type)
info.names[SELF_UVAR_NAME] = SymbolTableNode(MDEF, self_tvar_expr)
replace_tvar_def = TypeVarDef(SELF_UVAR_NAME,
info.fullname + '.' + SELF_UVAR_NAME,
-1, [], fill_typevars(info))
replace_other_type = TypeVarType(replace_tvar_def)
add_method(ctx,
'replace',
args=[
Argument(
Var('changes', AnyType(TypeOfAny.explicit)),
AnyType(TypeOfAny.explicit), None, ARG_STAR2)
],
return_type=replace_other_type,
self_type=replace_other_type,
tvar_def=replace_tvar_def)
# If there are no attributes, it may be that the semantic analyzer has not
# processed them yet. In order to work around this, we can simply skip generating
# __init__ if there are no attributes, because if the user truly did not define any,
# then the object default __init__ with an empty signature will be present anyway.
if (decorator_arguments['init']
and ('__init__' not in info.names
or info.names['__init__'].plugin_generated)
and attributes):
add_method(
ctx,
'__init__',
args=[
attr.to_argument() for attr in attributes
if attr.is_in_init
],
return_type=NoneType(),
)
if (decorator_arguments['eq'] and info.get('__eq__') is None
or decorator_arguments['order']):
# Type variable for self types in generated methods.
obj_type = ctx.api.named_type('__builtins__.object')
self_tvar_expr = TypeVarExpr(SELF_TVAR_NAME,
info.fullname + '.' + SELF_TVAR_NAME,
[], obj_type)
info.names[SELF_TVAR_NAME] = SymbolTableNode(MDEF, self_tvar_expr)
# Add <, >, <=, >=, but only if the class has an eq method.
if decorator_arguments['order']:
if not decorator_arguments['eq']:
ctx.api.fail('eq must be True if order is True', ctx.cls)
for method_name in ['__lt__', '__gt__', '__le__', '__ge__']:
# Like for __eq__ and __ne__, we want "other" to match
# the self type.
obj_type = ctx.api.named_type('__builtins__.object')
order_tvar_def = TypeVarDef(
SELF_TVAR_NAME, info.fullname + '.' + SELF_TVAR_NAME, -1,
[], obj_type)
order_other_type = TypeVarType(order_tvar_def)
order_return_type = ctx.api.named_type('__builtins__.bool')
order_args = [
Argument(Var('other', order_other_type), order_other_type,
None, ARG_POS)
]
existing_method = info.get(method_name)
if existing_method is not None and not existing_method.plugin_generated:
assert existing_method.node
ctx.api.fail(
'You may not have a custom %s method when order=True' %
method_name,
existing_method.node,
)
add_method(
ctx,
method_name,
args=order_args,
return_type=order_return_type,
self_type=order_other_type,
tvar_def=order_tvar_def,
)
if decorator_arguments['frozen']:
self._freeze(attributes)
self.reset_init_only_vars(info, attributes)
info.metadata['dataclass'] = {
'attributes': [attr.serialize() for attr in attributes],
'frozen': decorator_arguments['frozen'],
}
def build_namedtuple_typeinfo(self,
name: str,
items: List[str],
types: List[Type],
default_items: Mapping[str, Expression],
line: int) -> TypeInfo:
strtype = self.api.named_type('__builtins__.str')
implicit_any = AnyType(TypeOfAny.special_form)
basetuple_type = self.api.named_type('__builtins__.tuple', [implicit_any])
dictype = (self.api.named_type_or_none('builtins.dict', [strtype, implicit_any])
or self.api.named_type('__builtins__.object'))
# Actual signature should return OrderedDict[str, Union[types]]
ordereddictype = (self.api.named_type_or_none('builtins.dict', [strtype, implicit_any])
or self.api.named_type('__builtins__.object'))
fallback = self.api.named_type('__builtins__.tuple', [implicit_any])
# Note: actual signature should accept an invariant version of Iterable[UnionType[types]].
# but it can't be expressed. 'new' and 'len' should be callable types.
iterable_type = self.api.named_type_or_none('typing.Iterable', [implicit_any])
function_type = self.api.named_type('__builtins__.function')
info = self.api.basic_new_typeinfo(name, fallback)
info.is_named_tuple = True
tuple_base = TupleType(types, fallback)
info.tuple_type = tuple_base
info.line = line
# We can't calculate the complete fallback type until after semantic
# analysis, since otherwise base classes might be incomplete. Postpone a
# callback function that patches the fallback.
self.api.schedule_patch(PRIORITY_FALLBACKS,
lambda: calculate_tuple_fallback(tuple_base))
def add_field(var: Var, is_initialized_in_class: bool = False,
is_property: bool = False) -> None:
var.info = info
var.is_initialized_in_class = is_initialized_in_class
var.is_property = is_property
var._fullname = '%s.%s' % (info.fullname, var.name)
info.names[var.name] = SymbolTableNode(MDEF, var)
fields = [Var(item, typ) for item, typ in zip(items, types)]
for var in fields:
add_field(var, is_property=True)
# We can't share Vars between fields and method arguments, since they
# have different full names (the latter are normally used as local variables
# in functions, so their full names are set to short names when generated methods
# are analyzed).
vars = [Var(item, typ) for item, typ in zip(items, types)]
tuple_of_strings = TupleType([strtype for _ in items], basetuple_type)
add_field(Var('_fields', tuple_of_strings), is_initialized_in_class=True)
add_field(Var('_field_types', dictype), is_initialized_in_class=True)
add_field(Var('_field_defaults', dictype), is_initialized_in_class=True)
add_field(Var('_source', strtype), is_initialized_in_class=True)
add_field(Var('__annotations__', ordereddictype), is_initialized_in_class=True)
add_field(Var('__doc__', strtype), is_initialized_in_class=True)
tvd = TypeVarDef(SELF_TVAR_NAME, info.fullname + '.' + SELF_TVAR_NAME,
-1, [], info.tuple_type)
selftype = TypeVarType(tvd)
def add_method(funcname: str,
ret: Type,
args: List[Argument],
is_classmethod: bool = False,
is_new: bool = False,
) -> None:
if is_classmethod or is_new:
first = [Argument(Var('_cls'), TypeType.make_normalized(selftype), None, ARG_POS)]
else:
first = [Argument(Var('_self'), selftype, None, ARG_POS)]
args = first + args
types = [arg.type_annotation for arg in args]
items = [arg.variable.name for arg in args]
arg_kinds = [arg.kind for arg in args]
assert None not in types
signature = CallableType(cast(List[Type], types), arg_kinds, items, ret,
function_type)
signature.variables = [tvd]
func = FuncDef(funcname, args, Block([]))
func.info = info
func.is_class = is_classmethod
func.type = set_callable_name(signature, func)
func._fullname = info.fullname + '.' + funcname
func.line = line
if is_classmethod:
v = Var(funcname, func.type)
v.is_classmethod = True
v.info = info
v._fullname = func._fullname
func.is_decorated = True
dec = Decorator(func, [NameExpr('classmethod')], v)
dec.line = line
sym = SymbolTableNode(MDEF, dec)
else:
sym = SymbolTableNode(MDEF, func)
sym.plugin_generated = True
info.names[funcname] = sym
add_method('_replace', ret=selftype,
args=[Argument(var, var.type, EllipsisExpr(), ARG_NAMED_OPT) for var in vars])
def make_init_arg(var: Var) -> Argument:
default = default_items.get(var.name, None)
kind = ARG_POS if default is None else ARG_OPT
return Argument(var, var.type, default, kind)
add_method('__new__', ret=selftype,
args=[make_init_arg(var) for var in vars],
is_new=True)
add_method('_asdict', args=[], ret=ordereddictype)
special_form_any = AnyType(TypeOfAny.special_form)
add_method('_make', ret=selftype, is_classmethod=True,
args=[Argument(Var('iterable', iterable_type), iterable_type, None, ARG_POS),
Argument(Var('new'), special_form_any, EllipsisExpr(), ARG_NAMED_OPT),
Argument(Var('len'), special_form_any, EllipsisExpr(), ARG_NAMED_OPT)])
self_tvar_expr = TypeVarExpr(SELF_TVAR_NAME, info.fullname + '.' + SELF_TVAR_NAME,
[], info.tuple_type)
info.names[SELF_TVAR_NAME] = SymbolTableNode(MDEF, self_tvar_expr)
return info
def visit_type_var_expr(self, node: TypeVarExpr) -> Node:
return TypeVarExpr(node.name(), node.fullname(),
self.types(node.values))
def transform(self) -> None:
"""Apply all the necessary transformations to the underlying
dataclass so as to ensure it is fully type checked according
to the rules in PEP 557.
"""
ctx = self._ctx
info = self._ctx.cls.info
attributes = self.collect_attributes()
if attributes is None:
# Some definitions are not ready, defer() should be already called.
return
for attr in attributes:
if attr.type is None:
ctx.api.defer()
return
decorator_arguments = {
'init':
_get_decorator_bool_argument(self._ctx, 'init', True),
'eq':
_get_decorator_bool_argument(self._ctx, 'eq', True),
'order':
_get_decorator_bool_argument(self._ctx, 'order', False),
'frozen':
_get_decorator_bool_argument(self._ctx, 'frozen', False),
'slots':
_get_decorator_bool_argument(self._ctx, 'slots', False),
'match_args':
_get_decorator_bool_argument(self._ctx, 'match_args', True),
}
py_version = self._ctx.api.options.python_version
# If there are no attributes, it may be that the semantic analyzer has not
# processed them yet. In order to work around this, we can simply skip generating
# __init__ if there are no attributes, because if the user truly did not define any,
# then the object default __init__ with an empty signature will be present anyway.
if (decorator_arguments['init']
and ('__init__' not in info.names
or info.names['__init__'].plugin_generated)
and attributes):
args = [
attr.to_argument() for attr in attributes
if attr.is_in_init and not self._is_kw_only_type(attr.type)
]
if info.fallback_to_any:
# Make positional args optional since we don't know their order.
# This will at least allow us to typecheck them if they are called
# as kwargs
for arg in args:
if arg.kind == ARG_POS:
arg.kind = ARG_OPT
nameless_var = Var('')
args = [
Argument(nameless_var, AnyType(TypeOfAny.explicit), None,
ARG_STAR),
*args,
Argument(nameless_var, AnyType(TypeOfAny.explicit), None,
ARG_STAR2),
]
add_method(
ctx,
'__init__',
args=args,
return_type=NoneType(),
)
if (decorator_arguments['eq'] and info.get('__eq__') is None
or decorator_arguments['order']):
# Type variable for self types in generated methods.
obj_type = ctx.api.named_type('builtins.object')
self_tvar_expr = TypeVarExpr(SELF_TVAR_NAME,
info.fullname + '.' + SELF_TVAR_NAME,
[], obj_type)
info.names[SELF_TVAR_NAME] = SymbolTableNode(MDEF, self_tvar_expr)
# Add <, >, <=, >=, but only if the class has an eq method.
if decorator_arguments['order']:
if not decorator_arguments['eq']:
ctx.api.fail('eq must be True if order is True', ctx.cls)
for method_name in ['__lt__', '__gt__', '__le__', '__ge__']:
# Like for __eq__ and __ne__, we want "other" to match
# the self type.
obj_type = ctx.api.named_type('builtins.object')
order_tvar_def = TypeVarType(
SELF_TVAR_NAME, info.fullname + '.' + SELF_TVAR_NAME, -1,
[], obj_type)
order_return_type = ctx.api.named_type('builtins.bool')
order_args = [
Argument(Var('other', order_tvar_def), order_tvar_def,
None, ARG_POS)
]
existing_method = info.get(method_name)
if existing_method is not None and not existing_method.plugin_generated:
assert existing_method.node
ctx.api.fail(
'You may not have a custom %s method when order=True' %
method_name,
existing_method.node,
)
add_method(
ctx,
method_name,
args=order_args,
return_type=order_return_type,
self_type=order_tvar_def,
tvar_def=order_tvar_def,
)
if decorator_arguments['frozen']:
self._freeze(attributes)
else:
self._propertize_callables(attributes)
if decorator_arguments['slots']:
self.add_slots(info,
attributes,
correct_version=py_version >= (3, 10))
self.reset_init_only_vars(info, attributes)
if (decorator_arguments['match_args']
and ('__match_args__' not in info.names
or info.names['__match_args__'].plugin_generated)
and attributes):
str_type = ctx.api.named_type("builtins.str")
literals: List[Type] = [
LiteralType(attr.name, str_type) for attr in attributes
if attr.is_in_init
]
match_args_type = TupleType(literals,
ctx.api.named_type("builtins.tuple"))
add_attribute_to_class(ctx.api,
ctx.cls,
"__match_args__",
match_args_type,
final=True)
self._add_dataclass_fields_magic_attribute()
info.metadata['dataclass'] = {
'attributes': [attr.serialize() for attr in attributes],
'frozen': decorator_arguments['frozen'],
}
def visit_type_var_expr(self, node: TypeVarExpr) -> TypeVarExpr:
return TypeVarExpr(node.name, node.fullname,
self.types(node.values),
self.type(node.upper_bound), variance=node.variance)
def transform(self) -> None:
"""Apply all the necessary transformations to the underlying
dataclass so as to ensure it is fully type checked according
to the rules in PEP 557.
"""
ctx = self._ctx
info = self._ctx.cls.info
attributes = self.collect_attributes()
if attributes is None:
# Some definitions are not ready, defer() should be already called.
return
for attr in attributes:
if attr.type is None:
ctx.api.defer()
return
decorator_arguments = {
"init": _get_decorator_bool_argument(self._ctx, "init", True),
"eq": _get_decorator_bool_argument(self._ctx, "eq", True),
"order": _get_decorator_bool_argument(self._ctx, "order", False),
"frozen": _get_decorator_bool_argument(self._ctx, "frozen", False),
}
# If there are no attributes, it may be that the semantic analyzer has not
# processed them yet. In order to work around this, we can simply skip
# generating __init__ if there are no attributes, because if the user
# truly did not define any, then the object default __init__ with an
# empty signature will be present anyway.
if (decorator_arguments["init"]
and ("__init__" not in info.names
or info.names["__init__"].plugin_generated)
and attributes):
add_method(
ctx,
"__init__",
args=[
attr.to_argument() for attr in attributes
if attr.is_in_init
],
return_type=NoneType(),
)
if (decorator_arguments["eq"] and info.get("__eq__") is None
or decorator_arguments["order"]):
# Type variable for self types in generated methods.
obj_type = ctx.api.named_type("__builtins__.object")
self_tvar_expr = TypeVarExpr(SELF_TVAR_NAME,
info.fullname + "." + SELF_TVAR_NAME,
[], obj_type)
info.names[SELF_TVAR_NAME] = SymbolTableNode(MDEF, self_tvar_expr)
# Add <, >, <=, >=, but only if the class has an eq method.
if decorator_arguments["order"]:
if not decorator_arguments["eq"]:
ctx.api.fail("eq must be True if order is True", ctx.cls)
for method_name in ["__lt__", "__gt__", "__le__", "__ge__"]:
# Like for __eq__ and __ne__, we want "other" to match
# the self type.
obj_type = ctx.api.named_type("__builtins__.object")
order_tvar_def = TypeVarDef(
SELF_TVAR_NAME,
info.fullname + "." + SELF_TVAR_NAME,
-1,
[],
obj_type,
)
order_other_type = TypeVarType(order_tvar_def)
order_return_type = ctx.api.named_type("__builtins__.bool")
order_args = [
Argument(Var("other", order_other_type), order_other_type,
None, ARG_POS)
]
existing_method = info.get(method_name)
if existing_method is not None and not existing_method.plugin_generated:
assert existing_method.node
ctx.api.fail(
"You may not have a custom %s method when order=True" %
method_name,
existing_method.node,
)
add_method(
ctx,
method_name,
args=order_args,
return_type=order_return_type,
self_type=order_other_type,
tvar_def=order_tvar_def,
)
if decorator_arguments["frozen"]:
self._freeze(attributes)
self.reset_init_only_vars(info, attributes)
info.metadata["dataclass"] = {
"attributes": [attr.serialize() for attr in attributes],
"frozen": decorator_arguments["frozen"],
}
def transform(self) -> None:
"""Apply all the necessary transformations to the underlying
dataclass so as to ensure it is fully type checked according
to the rules in PEP 557.
"""
ctx = self._ctx
info = self._ctx.cls.info
attributes = self.collect_attributes()
if ctx.api.options.new_semantic_analyzer:
# Check if attribute types are ready.
for attr in attributes:
if info[attr.name].type is None:
ctx.api.defer()
return
decorator_arguments = {
'init': _get_decorator_bool_argument(self._ctx, 'init', True),
'eq': _get_decorator_bool_argument(self._ctx, 'eq', True),
'order': _get_decorator_bool_argument(self._ctx, 'order', False),
'frozen': _get_decorator_bool_argument(self._ctx, 'frozen', False),
}
if decorator_arguments['init']:
add_method(
ctx,
'__init__',
args=[attr.to_argument(info) for attr in attributes if attr.is_in_init],
return_type=NoneType(),
)
if (decorator_arguments['eq'] and info.get('__eq__') is None or
decorator_arguments['order']):
# Type variable for self types in generated methods.
obj_type = ctx.api.named_type('__builtins__.object')
self_tvar_expr = TypeVarExpr(SELF_TVAR_NAME, info.fullname() + '.' + SELF_TVAR_NAME,
[], obj_type)
info.names[SELF_TVAR_NAME] = SymbolTableNode(MDEF, self_tvar_expr)
# Add an eq method, but only if the class doesn't already have one.
if decorator_arguments['eq'] and info.get('__eq__') is None:
for method_name in ['__eq__', '__ne__']:
# The TVar is used to enforce that "other" must have
# the same type as self (covariant). Note the
# "self_type" parameter to add_method.
obj_type = ctx.api.named_type('__builtins__.object')
cmp_tvar_def = TypeVarDef(SELF_TVAR_NAME, info.fullname() + '.' + SELF_TVAR_NAME,
-1, [], obj_type)
cmp_other_type = TypeVarType(cmp_tvar_def)
cmp_return_type = ctx.api.named_type('__builtins__.bool')
add_method(
ctx,
method_name,
args=[Argument(Var('other', cmp_other_type), cmp_other_type, None, ARG_POS)],
return_type=cmp_return_type,
self_type=cmp_other_type,
tvar_def=cmp_tvar_def,
)
# Add <, >, <=, >=, but only if the class has an eq method.
if decorator_arguments['order']:
if not decorator_arguments['eq']:
ctx.api.fail('eq must be True if order is True', ctx.cls)
for method_name in ['__lt__', '__gt__', '__le__', '__ge__']:
# Like for __eq__ and __ne__, we want "other" to match
# the self type.
obj_type = ctx.api.named_type('__builtins__.object')
order_tvar_def = TypeVarDef(SELF_TVAR_NAME, info.fullname() + '.' + SELF_TVAR_NAME,
-1, [], obj_type)
order_other_type = TypeVarType(order_tvar_def)
order_return_type = ctx.api.named_type('__builtins__.bool')
order_args = [
Argument(Var('other', order_other_type), order_other_type, None, ARG_POS)
]
existing_method = info.get(method_name)
if existing_method is not None and not existing_method.plugin_generated:
assert existing_method.node
ctx.api.fail(
'You may not have a custom %s method when order=True' % method_name,
existing_method.node,
)
add_method(
ctx,
method_name,
args=order_args,
return_type=order_return_type,
self_type=order_other_type,
tvar_def=order_tvar_def,
)
if decorator_arguments['frozen']:
self._freeze(attributes)
self.reset_init_only_vars(info, attributes)
info.metadata['dataclass'] = {
'attributes': OrderedDict((attr.name, attr.serialize()) for attr in attributes),
'frozen': decorator_arguments['frozen'],
}
def visit_type_var_expr(self, node: TypeVarExpr) -> Node:
return TypeVarExpr(node.name(), node.fullname(),
self.types(node.values))
def _add_tvar_expr(self, name: str, ctx):
info = ctx.cls.info
obj_type = ctx.api.named_type("__builtins__.object")
self_tvar_expr = TypeVarExpr(name, self._get_tvar_name(name, info), [],
obj_type)
info.names[name] = SymbolTableNode(MDEF, self_tvar_expr)
