def test_generic_function_type(self):
c = Callable([self.x, self.y], [ARG_POS, ARG_POS], [None, None],
self.y, False, None, [TypeVarDef('X', -1, None)])
assert_equal(str(c), 'def [X] (X?, Y?) -> Y?')
v = [TypeVarDef('Y', -1, None), TypeVarDef('X', -2, None)]
c2 = Callable([], [], [], Void(None), False, None, v)
assert_equal(str(c2), 'def [Y, X] ()')
def test_generic_function_type(self) -> None:
c = CallableType([self.x, self.y], [ARG_POS, ARG_POS], [None, None],
self.y, self.function, name=None,
variables=[TypeVarDef('X', 'X', -1, [], self.fx.o)])
assert_equal(str(c), 'def [X] (X?, Y?) -> Y?')
v = [TypeVarDef('Y', 'Y', -1, [], self.fx.o),
TypeVarDef('X', 'X', -2, [], self.fx.o)]
c2 = CallableType([], [], [], NoneType(), self.function, name=None, variables=v)
assert_equal(str(c2), 'def [Y, X] ()')
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 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 make_type_info(name: str,
is_abstract: bool = False,
mro: List[TypeInfo] = None,
bases: List[Instance] = None,
typevars: List[str] = None) -> TypeInfo:
"""Make a TypeInfo suitable for use in unit tests."""
type_def = TypeDef(name, Block([]), None, [])
type_def.fullname = name
if typevars:
v = [] # type: List[TypeVarDef]
id = 1
for n in typevars:
v.append(TypeVarDef(n, id, None))
id += 1
type_def.type_vars = v
info = TypeInfo(SymbolTable(), type_def)
if mro is None:
mro = []
info.mro = [info] + mro
if bases is None:
if mro:
# By default, assume that there is a single non-generic base.
bases = [Instance(mro[0], [])]
else:
bases = []
info.bases = bases
return info
def add_class_tvars(t: Type, itype: Instance, is_classmethod: bool,
builtin_type: Callable[[str], Instance],
original_type: Type) -> Type:
"""Instantiate type variables during analyze_class_attribute_access,
e.g T and Q in the following:
def A(Generic(T)):
@classmethod
def foo(cls: Type[Q]) -> Tuple[T, Q]: ...
class B(A): pass
B.foo()
original_type is the value of the type B in the expression B.foo()
"""
# TODO: verify consistency between Q and T
info = itype.type # type: TypeInfo
if isinstance(t, CallableType):
# TODO: Should we propagate type variable values?
tvars = [TypeVarDef(n, n, i + 1, [], builtin_type('builtins.object'), tv.variance)
for (i, n), tv in zip(enumerate(info.type_vars), info.defn.type_vars)]
if is_classmethod:
t = bind_self(t, original_type, is_classmethod=True)
return t.copy_modified(variables=tvars + t.variables)
elif isinstance(t, Overloaded):
return Overloaded([cast(CallableType, add_class_tvars(item, itype, is_classmethod,
builtin_type, original_type))
for item in t.items()])
return t
def anal_var_defs(self, var_defs: List[TypeVarDef]) -> List[TypeVarDef]:
a = [] # type: List[TypeVarDef]
for vd in var_defs:
a.append(
TypeVarDef(vd.name, vd.id.raw_id, self.anal_array(vd.values),
vd.upper_bound.accept(self), vd.variance, vd.line))
return a
def add_class_tvars(t: Type, info: TypeInfo, is_classmethod: bool,
builtin_type: Callable[[str], Instance]) -> Type:
if isinstance(t, CallableType):
# TODO: Should we propagate type variable values?
vars = [
TypeVarDef(n, i + 1, None, builtin_type('builtins.object'),
tv.variance)
for (i,
n), tv in zip(enumerate(info.type_vars), info.defn.type_vars)
]
arg_types = t.arg_types
arg_kinds = t.arg_kinds
arg_names = t.arg_names
if is_classmethod:
arg_types = arg_types[1:]
arg_kinds = arg_kinds[1:]
arg_names = arg_names[1:]
return t.copy_modified(arg_types=arg_types,
arg_kinds=arg_kinds,
arg_names=arg_names,
variables=vars + t.variables)
elif isinstance(t, Overloaded):
return Overloaded([
cast(CallableType,
add_class_tvars(i, info, is_classmethod, builtin_type))
for i in t.items()
])
return t
def anal_var_defs(self, var_defs: List[TypeVarDef]) -> List[TypeVarDef]:
a = List[TypeVarDef]()
for vd in var_defs:
a.append(
TypeVarDef(vd.name, vd.id, self.anal_array(vd.values),
vd.upper_bound.accept(self), vd.line, vd.repr))
return a
def add_class_tvars(t: Type, info: TypeInfo, is_classmethod: bool,
builtin_type: Callable[[str], Instance]) -> Type:
if isinstance(t, CallableType):
# TODO: Should we propagate type variable values?
vars = [TypeVarDef(n, i + 1, None, builtin_type('builtins.object'))
for i, n in enumerate(info.type_vars)]
arg_types = t.arg_types
arg_kinds = t.arg_kinds
arg_names = t.arg_names
if is_classmethod:
arg_types = arg_types[1:]
arg_kinds = arg_kinds[1:]
arg_names = arg_names[1:]
return CallableType(arg_types,
arg_kinds,
arg_names,
t.ret_type,
t.fallback,
t.name,
vars + t.variables,
t.bound_vars,
t.line)
elif isinstance(t, Overloaded):
return Overloaded([cast(CallableType, add_class_tvars(i, info, is_classmethod, builtin_type))
for i in t.items()])
return t
def bind_new(self, name: str, tvar_expr: TypeVarLikeExpr) -> TypeVarLikeDef:
if self.is_class_scope:
self.class_id += 1
i = self.class_id
else:
self.func_id -= 1
i = self.func_id
if isinstance(tvar_expr, TypeVarExpr):
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
) # type: TypeVarLikeDef
elif isinstance(tvar_expr, ParamSpecExpr):
tvar_def = ParamSpecDef(
name,
tvar_expr.fullname,
i,
line=tvar_expr.line,
column=tvar_expr.column
)
else:
assert False
self.scope[tvar_expr.fullname] = tvar_def
return tvar_def
def add_class_tvars(t: Type, itype: Instance, isuper: Optional[Instance],
is_classmethod: bool, builtin_type: Callable[[str],
Instance],
original_type: Type) -> Type:
"""Instantiate type variables during analyze_class_attribute_access,
e.g T and Q in the following:
class A(Generic[T]):
@classmethod
def foo(cls: Type[Q]) -> Tuple[T, Q]: ...
class B(A[str]): pass
B.foo()
original_type is the value of the type B in the expression B.foo()
"""
# TODO: verify consistency between Q and T
info = itype.type # type: TypeInfo
if is_classmethod:
assert isuper is not None
t = expand_type_by_instance(t, isuper)
# We add class type variables if the class method is accessed on class object
# without applied type arguments, this matches the behavior of __init__().
# For example (continuing the example in docstring):
# A # The type of callable is def [T] () -> A[T], _not_ def () -> A[Any]
# A[int] # The type of callable is def () -> A[int]
# and
# A.foo # The type is generic def [T] () -> Tuple[T, A[T]]
# A[int].foo # The type is non-generic def () -> Tuple[int, A[int]]
#
# This behaviour is useful for defining alternative constructors for generic classes.
# To achieve such behaviour, we add the class type variables that are still free
# (i.e. appear in the return type of the class object on which the method was accessed).
free_ids = {t.id for t in itype.args if isinstance(t, TypeVarType)}
if isinstance(t, CallableType):
# NOTE: in practice either all or none of the variables are free, since
# visit_type_application() will detect any type argument count mismatch and apply
# a correct number of Anys.
tvars = [
TypeVarDef(n, n, i + 1, [], builtin_type('builtins.object'),
tv.variance)
for (i,
n), tv in zip(enumerate(info.type_vars), info.defn.type_vars)
# use 'is' to avoid id clashes with unrelated variables
if any(tv.id is id for id in free_ids)
]
if is_classmethod:
t = bind_self(t, original_type, is_classmethod=True)
return t.copy_modified(variables=tvars + t.variables)
elif isinstance(t, Overloaded):
return Overloaded([
cast(
CallableType,
add_class_tvars(item, itype, isuper, is_classmethod,
builtin_type, original_type))
for item in t.items()
])
return t
def _type_var_def(name: str,
module: str,
upper_bound,
values=(),
meta_level=0) -> TypeVarDef:
id_ = TypeVarId.new(meta_level)
id_.raw_id = -id_.raw_id
fullname = f'{module}.{name}'
return TypeVarDef(name, fullname, id_, list(values), upper_bound)
def translate_variables(self,
variables: List[TypeVarDef]) -> List[TypeVarDef]:
if not variables:
return variables
items = [] # type: List[TypeVarDef]
for v in variables:
if v.id > 0:
items.append(
TypeVarDef(v.name, -v.id - self.num_func_tvars, v.values))
else:
items.append(v)
return items
def callable(self, vars, *a) -> Callable:
"""callable(args, a1, ..., an, r) constructs a callable with
argument types a1, ... an and return type r and type arguments
vars.
"""
tv = [] # type: List[TypeVarDef]
n = -1
for v in vars:
tv.append(TypeVarDef(v, n, None))
n -= 1
return Callable(a[:-1], [ARG_POS] * (len(a) - 1),
[None] * (len(a) - 1), a[-1], False, None, tv)
def add_class_tvars(t: Type, info: TypeInfo) -> Type:
if isinstance(t, Callable):
vars = [
TypeVarDef(n, i + 1, None) for i, n in enumerate(info.type_vars)
]
return Callable(t.arg_types, t.arg_kinds, t.arg_names, t.ret_type,
t.is_type_obj(), t.name, vars + t.variables,
t.bound_vars, t.line, None)
elif isinstance(t, Overloaded):
return Overloaded(
[cast(Callable, add_class_tvars(i, info)) for i in t.items()])
return t
def class_callable(init_type: Callable, info: TypeInfo) -> Callable:
"""Create a type object type based on the signature of __init__."""
variables = [] # type: List[TypeVarDef]
for i, tvar in enumerate(info.defn.type_vars):
variables.append(TypeVarDef(tvar.name, i + 1, tvar.values))
initvars = init_type.variables
variables.extend(initvars)
c = Callable(init_type.arg_types, init_type.arg_kinds, init_type.arg_names,
self_type(info), True, None,
variables).with_name('"{}"'.format(info.name()))
return convert_class_tvars_to_func_tvars(c, len(initvars))
def bind(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, 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 class_callable(init_type: CallableType, info: TypeInfo, type_type: Instance) -> CallableType:
"""Create a type object type based on the signature of __init__."""
variables = [] # type: List[TypeVarDef]
for i, tvar in enumerate(info.defn.type_vars):
variables.append(TypeVarDef(tvar.name, i + 1, tvar.values, tvar.upper_bound,
tvar.variance))
initvars = init_type.variables
variables.extend(initvars)
callable_type = init_type.copy_modified(
ret_type=self_type(info), fallback=type_type, name=None, variables=variables)
c = callable_type.with_name('"{}"'.format(info.name()))
return convert_class_tvars_to_func_tvars(c, len(initvars))
def callable(self, vars: List[str], *a: Type) -> CallableType:
"""callable(args, a1, ..., an, r) constructs a callable with
argument types a1, ... an and return type r and type arguments
vars.
"""
tv = [] # type: List[TypeVarDef]
n = -1
for v in vars:
tv.append(TypeVarDef(v, v, n, [], self.fx.o))
n -= 1
return CallableType(list(a[:-1]), [ARG_POS] * (len(a) - 1),
[None] * (len(a) - 1),
a[-1],
self.fx.function,
name=None,
variables=tv)
def freshen_function_type_vars(callee: F) -> F:
"""Substitute fresh type variables for generic function type variables."""
if isinstance(callee, CallableType):
if not callee.is_generic():
return cast(F, callee)
tvdefs = []
tvmap = {} # type: Dict[TypeVarId, Type]
for v in callee.variables:
tvdef = TypeVarDef.new_unification_variable(v)
tvdefs.append(tvdef)
tvmap[v.id] = TypeVarType(tvdef)
fresh = cast(CallableType, expand_type(callee, tvmap)).copy_modified(variables=tvdefs)
return cast(F, fresh)
else:
assert isinstance(callee, Overloaded)
fresh_overload = Overloaded([freshen_function_type_vars(item)
for item in callee.items()])
return cast(F, fresh_overload)
def freshen_function_type_vars(callee: F) -> F:
"""Substitute fresh type variables for generic function type variables."""
if isinstance(callee, CallableType):
if not callee.is_generic():
return cast(F, callee)
tvdefs = []
tvmap = {} # type: Dict[TypeVarId, Type]
for v in callee.variables:
tvdef = TypeVarDef.new_unification_variable(v)
tvdefs.append(tvdef)
tvmap[v.id] = TypeVarType(tvdef)
fresh = cast(CallableType, expand_type(callee, tvmap)).copy_modified(variables=tvdefs)
return cast(F, fresh)
else:
assert isinstance(callee, Overloaded)
fresh_overload = Overloaded([freshen_function_type_vars(item)
for item in callee.items()])
return cast(F, fresh_overload)
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('AT', 'AT', 1, [], object_type)
tvd_type = TypeVarType(tvd)
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 make_type_info(self,
name: str,
module_name: str = None,
is_abstract: bool = False,
mro: List[TypeInfo] = None,
bases: List[Instance] = None,
typevars: List[str] = None,
variances: List[int] = None) -> TypeInfo:
"""Make a TypeInfo suitable for use in unit tests."""
class_def = ClassDef(name, Block([]), None, [])
class_def.fullname = name
if module_name is None:
if '.' in name:
module_name = name.rsplit('.', 1)[0]
else:
module_name = '__main__'
if typevars:
v = [] # type: List[TypeVarDef]
for id, n in enumerate(typevars, 1):
if variances:
variance = variances[id - 1]
else:
variance = COVARIANT
v.append(TypeVarDef(n, id, None, self.o, variance=variance))
class_def.type_vars = v
info = TypeInfo(SymbolTable(), class_def, module_name)
if mro is None:
mro = []
if name != 'builtins.object':
mro.append(self.oi)
info.mro = [info] + mro
if bases is None:
if mro:
# By default, assume that there is a single non-generic base.
bases = [Instance(mro[0], [])]
else:
bases = []
info.bases = bases
return info
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 test_type_variable_binding(self) -> None:
assert_equal(str(TypeVarDef('X', 'X', 1, [], self.fx.o)), 'X')
assert_equal(str(TypeVarDef('X', 'X', 1, [self.x, self.y], self.fx.o)),
'X in (X?, Y?)')
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 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()
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=NoneTyp(),
)
for stmt in self._ctx.cls.defs.body:
# Fix up the types of classmethods since, by default,
# they will be based on the parent class' init.
if isinstance(stmt, Decorator) and stmt.func.is_class:
func_type = stmt.func.type
if isinstance(func_type, CallableType):
func_type.arg_types[0] = self._ctx.api.class_type(
self._ctx.cls.info)
if isinstance(stmt, OverloadedFuncDef) and stmt.is_class:
func_type = stmt.type
if isinstance(func_type, Overloaded):
class_type = ctx.api.class_type(ctx.cls.info)
for item in func_type.items():
item.arg_types[0] = class_type
if stmt.impl is not None:
assert isinstance(stmt.impl, Decorator)
if isinstance(stmt.impl.func.type, CallableType):
stmt.impl.func.type.arg_types[0] = class_type
# 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('T', 'T', -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('T', 'T', -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:
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)
# Remove init-only vars from the class.
for attr in attributes:
if attr.is_init_var:
del info.names[attr.name]
info.metadata['dataclass'] = {
'attributes':
OrderedDict((attr.name, 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 make_type_var(name: str, id: int, values: List[Type],
upper_bound: Type, variance: int) -> TypeVarType:
return TypeVarType(
TypeVarDef(name, id, values, upper_bound, 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"],
}
