def run_with_model_cls(self, model_cls: Type[Model]) -> None:
# get_FOO_display for choices
for field in self.django_context.get_model_fields(model_cls):
if field.choices:
info = self.lookup_typeinfo_or_incomplete_defn_error('builtins.str')
return_type = Instance(info, [])
common.add_method(self.ctx,
name='get_{}_display'.format(field.attname),
args=[],
return_type=return_type)
# get_next_by, get_previous_by for Date, DateTime
for field in self.django_context.get_model_fields(model_cls):
if isinstance(field, (DateField, DateTimeField)) and not field.null:
return_type = Instance(self.model_classdef.info, [])
common.add_method(self.ctx,
name='get_next_by_{}'.format(field.attname),
args=[Argument(Var('kwargs', AnyType(TypeOfAny.explicit)),
AnyType(TypeOfAny.explicit),
initializer=None,
kind=ARG_STAR2)],
return_type=return_type)
common.add_method(self.ctx,
name='get_previous_by_{}'.format(field.attname),
args=[Argument(Var('kwargs', AnyType(TypeOfAny.explicit)),
AnyType(TypeOfAny.explicit),
initializer=None,
kind=ARG_STAR2)],
return_type=return_type)
def add_method(funcname: str,
ret: Type,
args: List[Argument],
name: Optional[str] = None,
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
if is_classmethod:
v = Var(funcname, func.type)
v.is_classmethod = True
v.info = info
v._fullname = func._fullname
dec = Decorator(func, [NameExpr('classmethod')], v)
info.names[funcname] = SymbolTableNode(MDEF, dec)
else:
info.names[funcname] = SymbolTableNode(MDEF, func)
def copy_argument(self, argument: Argument) -> Argument:
init_stmt = None # type: AssignmentStmt
if argument.initialization_statement:
init_lvalue = cast(
NameExpr,
self.node(argument.initialization_statement.lvalues[0]),
)
init_lvalue.set_line(argument.line)
init_stmt = AssignmentStmt(
[init_lvalue],
self.node(argument.initialization_statement.rvalue),
self.optional_type(argument.initialization_statement.type),
)
arg = Argument(
self.visit_var(argument.variable),
argument.type_annotation,
argument.initializer,
argument.kind,
init_stmt,
)
# Refresh lines of the inner things
arg.set_line(argument.line)
return arg
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 copy_argument(self, argument: Argument) -> Argument:
init_stmt = None # type: AssignmentStmt
if argument.initialization_statement:
init_lvalue = cast(
NameExpr,
self.node(argument.initialization_statement.lvalues[0]),
)
init_lvalue.set_line(argument.line)
init_stmt = AssignmentStmt(
[init_lvalue],
self.node(argument.initialization_statement.rvalue),
self.optional_type(argument.initialization_statement.type),
)
arg = Argument(
self.visit_var(argument.variable),
argument.type_annotation,
argument.initializer,
argument.kind,
init_stmt,
)
# Refresh lines of the inner things
arg.set_line(argument.line)
return arg
def copy_method_to_another_class(
ctx: ClassDefContext, self_type: Instance, new_method_name: str, method_node: FuncDef
) -> None:
semanal_api = get_semanal_api(ctx)
if method_node.type is None:
if not semanal_api.final_iteration:
semanal_api.defer()
return
arguments, return_type = build_unannotated_method_args(method_node)
add_method(ctx, new_method_name, args=arguments, return_type=return_type, self_type=self_type)
return
method_type = method_node.type
if not isinstance(method_type, CallableType):
if not semanal_api.final_iteration:
semanal_api.defer()
return
arguments = []
bound_return_type = semanal_api.anal_type(method_type.ret_type, allow_placeholder=True)
assert bound_return_type is not None
if isinstance(bound_return_type, PlaceholderNode):
return
try:
original_arguments = method_node.arguments[1:]
except AttributeError:
original_arguments = []
for arg_name, arg_type, original_argument in zip(
method_type.arg_names[1:], method_type.arg_types[1:], original_arguments
):
bound_arg_type = semanal_api.anal_type(arg_type, allow_placeholder=True)
if bound_arg_type is None and not semanal_api.final_iteration:
semanal_api.defer()
return
assert bound_arg_type is not None
if isinstance(bound_arg_type, PlaceholderNode):
return
var = Var(name=original_argument.variable.name, type=arg_type)
var.line = original_argument.variable.line
var.column = original_argument.variable.column
argument = Argument(
variable=var,
type_annotation=bound_arg_type,
initializer=original_argument.initializer,
kind=original_argument.kind,
)
argument.set_line(original_argument)
arguments.append(argument)
add_method(ctx, new_method_name, args=arguments, return_type=bound_return_type, self_type=self_type)
def transform_args(
self,
n: ast27.arguments,
line: int,
) -> Tuple[List[Argument], List[Statement]]:
type_comments: Sequence[Optional[str]] = n.type_comments
converter = TypeConverter(self.errors,
line=line,
assume_str_is_unicode=self.unicode_literals)
decompose_stmts: List[Statement] = []
n_args = n.args
args = [(self.convert_arg(i, arg, line, decompose_stmts),
self.get_type(i, type_comments, converter))
for i, arg in enumerate(n_args)]
defaults = self.translate_expr_list(n.defaults)
names: List[str] = [
name for arg in n_args for name in self.extract_names(arg)
]
new_args: List[Argument] = []
num_no_defaults = len(args) - len(defaults)
# positional arguments without defaults
for a, annotation in args[:num_no_defaults]:
new_args.append(Argument(a, annotation, None, ARG_POS))
# positional arguments with defaults
for (a, annotation), d in zip(args[num_no_defaults:], defaults):
new_args.append(Argument(a, annotation, d, ARG_OPT))
# *arg
if n.vararg is not None:
new_args.append(
Argument(Var(n.vararg),
self.get_type(len(args), type_comments, converter),
None, ARG_STAR))
names.append(n.vararg)
# **kwarg
if n.kwarg is not None:
typ = self.get_type(
len(args) + (0 if n.vararg is None else 1), type_comments,
converter)
new_args.append(Argument(Var(n.kwarg), typ, None, ARG_STAR2))
names.append(n.kwarg)
for arg in new_args:
if argument_elide_name(arg.variable.name):
arg.pos_only = True
# We don't have any context object to give, but we have closed around the line num
def fail_arg(msg: str, arg: None) -> None:
self.fail(msg, line, 0)
check_arg_names(names, [None] * len(names), fail_arg)
return new_args, decompose_stmts
def copy_argument(self, argument: Argument) -> Argument:
arg = Argument(
self.visit_var(argument.variable),
argument.type_annotation,
argument.initializer,
argument.kind,
)
# Refresh lines of the inner things
arg.set_line(argument.line)
return arg
def copy_argument(self, argument: Argument) -> Argument:
arg = Argument(
self.visit_var(argument.variable),
argument.type_annotation,
argument.initializer,
argument.kind,
)
# Refresh lines of the inner things
arg.set_line(argument.line)
return arg
def transform_args(self, n: ast27.arguments, line: int) -> List[Argument]:
# TODO: remove the cast once https://github.com/python/typeshed/pull/522
# is accepted and synced
type_comments = cast(List[str], n.type_comments) # type: ignore
converter = TypeConverter(line=line)
def convert_arg(arg: ast27.expr) -> Var:
if isinstance(arg, ast27.Name):
v = arg.id
elif isinstance(arg, ast27.Tuple):
# TODO: An `arg` object may be a Tuple instead of just an identifier in the
# case of Python 2 function definitions/lambdas that use the tuple unpacking
# syntax. The `typed_ast.conversions` module ended up just simply passing the
# the arg object unmodified (instead of converting it into more args, etc).
# This isn't typesafe, since we will no longer be always passing in a string
# to `Var`, but we'll do the same here for consistency.
v = arg # type: ignore
else:
raise RuntimeError("'{}' is not a valid argument.".format(
ast27.dump(arg)))
return Var(v)
def get_type(i: int) -> Optional[Type]:
if i < len(type_comments) and type_comments[i] is not None:
return converter.visit_raw_str(type_comments[i])
return None
args = [(convert_arg(arg), get_type(i))
for i, arg in enumerate(n.args)]
defaults = self.translate_expr_list(n.defaults)
new_args = [] # type: List[Argument]
num_no_defaults = len(args) - len(defaults)
# positional arguments without defaults
for a, annotation in args[:num_no_defaults]:
new_args.append(Argument(a, annotation, None, ARG_POS))
# positional arguments with defaults
for (a, annotation), d in zip(args[num_no_defaults:], defaults):
new_args.append(Argument(a, annotation, d, ARG_OPT))
# *arg
if n.vararg is not None:
new_args.append(
Argument(Var(n.vararg), get_type(len(args)), None, ARG_STAR))
# **kwarg
if n.kwarg is not None:
typ = get_type(len(args) + (0 if n.vararg is None else 1))
new_args.append(Argument(Var(n.kwarg), typ, None, ARG_STAR2))
return new_args
def add_struc_and_unstruc_to_classdefcontext(cls_def_ctx: ClassDefContext):
"""This MyPy hook tells MyPy that struc and unstruc will be present on a Cat"""
dict_type = cls_def_ctx.api.named_type("__builtins__.dict")
str_type = cls_def_ctx.api.named_type("__builtins__.str")
api = cls_def_ctx.api
implicit_any = AnyType(TypeOfAny.special_form)
mapping = api.lookup_fully_qualified_or_none("typing.Mapping")
if not mapping or not mapping.node:
api.defer()
return
mapping_str_any_type = Instance(mapping.node, [str_type, implicit_any])
maybe_mapping_str_any_type = make_optional(mapping_str_any_type)
if fullname == CAT_PATH:
attr_class_maker_callback(
cls_def_ctx, True
) # since a Cat is also an attr.s class...
info = cls_def_ctx.cls.info
if STRUCTURE_NAME not in info.names:
add_static_method(
cls_def_ctx,
STRUCTURE_NAME,
[
Argument(
Var("d", mapping_str_any_type),
mapping_str_any_type,
None,
ARG_POS,
)
],
fill_typevars(info),
)
if TRY_STRUCTURE_NAME not in info.names:
# print('adding ' + TRY_STRUCTURE_NAME + ' to ' + str(info.fullname()) )
add_static_method(
cls_def_ctx,
TRY_STRUCTURE_NAME,
[
Argument(
Var("d", maybe_mapping_str_any_type),
maybe_mapping_str_any_type,
None,
ARG_POS,
)
],
fill_typevars(info),
)
if UNSTRUCTURE_NAME not in info.names:
add_method(cls_def_ctx, UNSTRUCTURE_NAME, [], dict_type)
def attr_utils_serialise_serde(cls_def_ctx: ClassDefContext):
"""
Handles :func:`attr_utils.serialise.serde`.
:param cls_def_ctx:
"""
info = cls_def_ctx.cls.info
# https://gitter.im/python/typing?at=5e078653eac8d1511e737d8c
str_type = cls_def_ctx.api.named_type(f"{_builtins}.str")
bool_type = cls_def_ctx.api.named_type(f"{_builtins}.bool")
implicit_any = AnyType(TypeOfAny.special_form)
mapping = cls_def_ctx.api.lookup_fully_qualified_or_none("typing.Mapping")
mutable_mapping = cls_def_ctx.api.lookup_fully_qualified_or_none(
"typing.MutableMapping")
mapping_str_any_type = Instance(mapping.node,
[str_type, implicit_any]) # type: ignore
mutable_mapping_str_any_type = Instance(
mutable_mapping.node, [str_type, implicit_any]) # type: ignore
# # maybe_mapping_str_any = UnionType.make_union([typ, NoneType()])(mapping_str_any_type)
decorated_class_instance = Instance(
cls_def_ctx.api.lookup_fully_qualified_or_none(
cls_def_ctx.cls.fullname).node, # type: ignore
[],
)
if "to_dict" not in info.names:
add_method_to_class(
api=cls_def_ctx.api,
cls=cls_def_ctx.cls,
name="to_dict",
args=[
Argument(Var("convert_values", bool_type), bool_type, None,
ARG_OPT)
],
return_type=mutable_mapping_str_any_type,
)
if "from_dict" not in info.names:
add_classmethod_to_class(
api=cls_def_ctx.api,
cls=cls_def_ctx.cls,
name="from_dict",
args=[
Argument(Var('d', mapping_str_any_type), mapping_str_any_type,
None, ARG_POS)
],
return_type=decorated_class_instance,
cls_type=TypeType(decorated_class_instance),
)
def add_get_set_attr_fallback_to_any(ctx: ClassDefContext):
any = AnyType(TypeOfAny.special_form)
name_arg = Argument(variable=Var('name', any),
type_annotation=any,
initializer=None,
kind=ARG_POS)
add_method(ctx, '__getattr__', [name_arg], any)
value_arg = Argument(variable=Var('value', any),
type_annotation=any,
initializer=None,
kind=ARG_POS)
add_method(ctx, '__setattr__', [name_arg, value_arg], any)
def strawberry_pydantic_class_callback(ctx: ClassDefContext) -> None:
# in future we want to have a proper pydantic plugin, but for now
# let's fallback to **kwargs for __init__, some resources are here:
# https://github.com/samuelcolvin/pydantic/blob/master/pydantic/mypy.py
# >>> model_index = ctx.cls.decorators[0].arg_names.index("model")
# >>> model_name = ctx.cls.decorators[0].args[model_index].name
# >>> model_type = ctx.api.named_type("UserModel")
# >>> model_type = ctx.api.lookup(model_name, Context())
model_expression = _get_argument(call=ctx.reason,
name="model") # type: ignore
if model_expression is None:
ctx.api.fail("model argument in decorator failed to be parsed",
ctx.reason)
else:
# Add __init__
init_args = [
Argument(Var("kwargs"), AnyType(TypeOfAny.explicit), None,
ARG_STAR2)
]
add_method(ctx, "__init__", init_args, NoneType())
model_type = _get_type_for_expr(model_expression, ctx.api)
# Add to_pydantic
add_method(
ctx,
"to_pydantic",
args=[],
return_type=model_type,
)
# Add from_pydantic
model_argument = Argument(
variable=Var(name="instance", type=model_type),
type_annotation=model_type,
initializer=None,
kind=ARG_OPT,
)
add_static_method_to_class(
ctx.api,
ctx.cls,
name="from_pydantic",
args=[model_argument],
return_type=fill_typevars(ctx.cls.info),
)
def add_dummy_init_method(ctx: ClassDefContext) -> None:
any = AnyType(TypeOfAny.special_form)
pos_arg = Argument(variable=Var('args', any),
type_annotation=any,
initializer=None,
kind=ARG_STAR)
kw_arg = Argument(variable=Var('kwargs', any),
type_annotation=any,
initializer=None,
kind=ARG_STAR2)
add_method(ctx, '__init__', [pos_arg, kw_arg], NoneTyp())
# mark as model class
ctx.cls.info.metadata.setdefault('django', {})['generated_init'] = True
def add_method(self,
method_name: str,
args: List[Argument],
ret_type: Type,
self_type: Optional[Type] = None,
tvd: Optional[TypeVarDef] = None) -> None:
"""Add a method: def <method_name>(self, <args>) -> <ret_type>): ... to info.
self_type: The type to use for the self argument or None to use the inferred self type.
tvd: If the method is generic these should be the type variables.
"""
from mypy.semanal import set_callable_name
self_type = self_type if self_type is not None else self.self_type
args = [Argument(Var('self'), self_type, None, ARG_POS)] + args
arg_types = [arg.type_annotation for arg in args]
arg_names = [arg.variable.name() for arg in args]
arg_kinds = [arg.kind for arg in args]
assert None not in arg_types
signature = CallableType(cast(List[Type], arg_types), arg_kinds,
arg_names, ret_type, self.function_type)
if tvd:
signature.variables = [tvd]
func = FuncDef(method_name, args, Block([PassStmt()]))
func.info = self.info
func.type = set_callable_name(signature, func)
func._fullname = self.info.fullname() + '.' + method_name
func.line = self.info.line
self.info.names[method_name] = SymbolTableNode(MDEF, func)
# Add the created methods to the body so that they can get further semantic analysis.
# e.g. Forward Reference Resolution.
self.info.defn.defs.body.append(func)
def add_classmethod_to_class(api,
cls,
name,
args,
return_type,
self_type=None,
tvar_def=None):
"""Adds a new classmethod to a class definition."""
info = cls.info
# First remove any previously generated methods with the same name
# to avoid clashes and problems in the semantic analyzer.
if name in info.names:
sym = info.names[name]
if sym.plugin_generated and isinstance(sym.node, Decorator):
cls.defs.body.remove(sym.node)
self_type = self_type or fill_typevars(info)
class_type = api.class_type(self_type)
function_type = builtin_type(api, 'function')
args = [Argument(Var('cls'), class_type, None, ARG_POS)] + args
arg_types, arg_names, arg_kinds = [], [], []
for arg in args:
assert arg.type_annotation, 'All arguments must be fully typed.'
arg_types.append(arg.type_annotation)
arg_names.append(arg.variable.name)
arg_kinds.append(arg.kind)
signature = CallableType(arg_types, arg_kinds, arg_names, return_type,
function_type)
if tvar_def:
signature.variables = [tvar_def]
func = FuncDef(name, args, Block([PassStmt()]))
func.type = set_callable_name(signature, func)
func._fullname = info.fullname + '.' + name
func.line = info.line
func.is_class = True
var = Var(name)
var.line = info.line
var.info = info
var.is_classmethod = True
# should we have a NameExpr in the decorator list?
dec = Decorator(func, [], var)
dec.line = info.line
# NOTE: we would like the plugin generated node to dominate, but we still
# need to keep any existing definitions so they get semantically analyzed.
if name in info.names:
# Get a nice unique name instead.
r_name = get_unique_redefinition_name(name, info.names)
info.names[r_name] = info.names[name]
info.names[name] = SymbolTableNode(MDEF, dec, plugin_generated=True)
info.defn.defs.body.append(dec)
def _adjust_interface_function(
self,
api: SemanticAnalyzerPluginInterface,
class_info: TypeInfo,
func_def: FuncDef,
) -> Statement:
if func_def.arg_names and func_def.arg_names[0] == "self":
# reveal the common mistake of leaving "self" arguments in the
# interface
api.fail("Interface methods should not have 'self' argument",
func_def)
else:
selftype = Instance(class_info, [],
line=class_info.line,
column=class_info.column)
selfarg = Argument(Var("self", None), selftype, None, ARG_POS)
if isinstance(func_def.type, CallableType):
func_def.type.arg_names.insert(0, "self")
func_def.type.arg_kinds.insert(0, ARG_POS)
func_def.type.arg_types.insert(0, selftype)
func_def.arg_names.insert(0, "self")
func_def.arg_kinds.insert(0, ARG_POS)
func_def.arguments.insert(0, selfarg)
return func_def
def add_model_init_hook(ctx: ClassDefContext) -> None:
"""Add a dummy __init__() to a model and record it is generated.
Instantiation will be checked more precisely when we inferred types
(using get_function_hook and model_hook).
"""
if '__init__' in ctx.cls.info.names:
# Don't override existing definition.
return
any = AnyType(TypeOfAny.special_form)
var = Var('kwargs', any)
kw_arg = Argument(variable=var,
type_annotation=any,
initializer=None,
kind=ARG_STAR2)
add_method(ctx, '__init__', [kw_arg], NoneTyp())
ctx.cls.info.metadata.setdefault('sqlalchemy', {})['generated_init'] = True
# Also add a selection of auto-generated attributes.
sym = ctx.api.lookup_fully_qualified_or_none('sqlalchemy.sql.schema.Table')
if sym:
assert isinstance(sym.node, TypeInfo)
typ = Instance(sym.node, []) # type: Type
else:
typ = AnyType(TypeOfAny.special_form)
add_var_to_class('__table__', typ, ctx.cls.info)
def argument(self, ctx: 'mypy.plugin.ClassDefContext') -> Argument:
"""Return this attribute as an argument to __init__."""
assert self.init
init_type = self.info[self.name].type
if self.converter_name:
# When a converter is set the init_type is overriden by the first argument
# of the converter method.
converter = ctx.api.lookup_fully_qualified(self.converter_name)
if (converter and converter.type
and isinstance(converter.type, CallableType)
and converter.type.arg_types):
init_type = converter.type.arg_types[0]
else:
init_type = None
if init_type is None:
if ctx.api.options.disallow_untyped_defs:
# This is a compromise. If you don't have a type here then the
# __init__ will be untyped. But since the __init__ is added it's
# pointing at the decorator. So instead we also show the error in the
# assignment, which is where you would fix the issue.
node = self.info[self.name].node
assert node is not None
ctx.api.msg.need_annotation_for_var(node, self.context)
# Convert type not set to Any.
init_type = AnyType(TypeOfAny.unannotated)
# Attrs removes leading underscores when creating the __init__ arguments.
return Argument(Var(self.name.lstrip("_"), init_type), init_type, None,
ARG_OPT if self.has_default else ARG_POS)
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_initializer(
self,
fields: List["PydanticModelField"],
config: "ModelConfigData",
is_settings: bool,
) -> None:
"""
Adds a fields-aware `__init__` method to the class.
The added `__init__` will be annotated with types vs. all `Any` depending on the plugin settings.
"""
ctx = self._ctx
typed = self.plugin_config.init_typed
use_alias = config.allow_population_by_field_name is not True
force_all_optional = is_settings or bool(
config.has_alias_generator
and not config.allow_population_by_field_name)
init_arguments = self.get_field_arguments(
fields,
typed=typed,
force_all_optional=force_all_optional,
use_alias=use_alias,
)
if not self.should_init_forbid_extra(fields, config):
var = Var("kwargs")
init_arguments.append(
Argument(var, AnyType(TypeOfAny.explicit), None, ARG_STAR2))
add_method(ctx, "__init__", init_arguments, NoneType())
def adjust_class_def(class_def_context: ClassDefContext) -> None:
# This MyPy plugin inserts method type stubs for the "missing" ordering methods the
# @total_ordering class decorator will fill in dynamically.
api = class_def_context.api
ordering_other_type = api.named_type("__builtins__.object")
ordering_return_type = api.named_type("__builtins__.bool")
args = [
Argument(
variable=Var(name="other", type=ordering_other_type),
type_annotation=ordering_other_type,
initializer=None,
kind=ARG_POS,
)
]
type_info: TypeInfo = class_def_context.cls.info
for ordering_method_name in "__lt__", "__le__", "__gt__", "__ge__":
existing_method = type_info.get(ordering_method_name)
if existing_method is None:
add_method(
ctx=class_def_context,
name=ordering_method_name,
args=args,
return_type=ordering_return_type,
)
def add_method(
ctx: ClassDefContext,
name: str,
args: List[Argument],
return_type: Type,
self_type: Optional[Type] = None,
tvar_def: Optional[TypeVarDef] = None,
) -> None:
"""Adds a new method to a class.
"""
info = ctx.cls.info
self_type = self_type or fill_typevars(info)
function_type = ctx.api.named_type('__builtins__.function')
args = [Argument(Var('self'), self_type, None, ARG_POS)] + args
arg_types, arg_names, arg_kinds = [], [], []
for arg in args:
assert arg.type_annotation, 'All arguments must be fully typed.'
arg_types.append(arg.type_annotation)
arg_names.append(arg.variable.name())
arg_kinds.append(arg.kind)
signature = CallableType(arg_types, arg_kinds, arg_names, return_type, function_type)
if tvar_def:
signature.variables = [tvar_def]
func = FuncDef(name, args, Block([PassStmt()]))
func.info = info
func.type = set_callable_name(signature, func)
func._fullname = info.fullname() + '.' + name
func.line = info.line
info.names[name] = SymbolTableNode(MDEF, func, plugin_generated=True)
info.defn.defs.body.append(func)
def argument(self) -> Argument:
"""Return this attribute as an argument to __init__."""
# Convert type not set to Any.
_type = self.type or AnyType(TypeOfAny.unannotated)
# Attrs removes leading underscores when creating the __init__ arguments.
return Argument(Var(self.name.lstrip("_"), _type), _type, None,
ARG_OPT if self.has_default else ARG_POS)
def constructor_args(self) -> List[Argument]:
return [
Argument(variable=Var(f'_{i}', t),
type_annotation=t,
initializer=None,
kind=ARG_POS) for i, t in enumerate(self.types)
]
def to_argument(self) -> Argument:
return Argument(
variable=self.to_var(),
type_annotation=self.type,
initializer=None,
kind=ARG_OPT if self.has_default else ARG_POS,
)
def to_argument(self, info: TypeInfo) -> Argument:
return Argument(
variable=self.to_var(info),
type_annotation=info[self.name].type,
initializer=None,
kind=ARG_OPT if self.has_default else ARG_POS,
)
def metharg(self):
kind = ARG_POS if not self.has_default else ARG_OPT
return Argument(
variable=self.var,
type_annotation=self.type,
initializer=None,
kind=kind,
)
def make_argument(arg: ast35.arg, default: Optional[ast35.expr],
kind: int) -> Argument:
if no_type_check:
arg_type = None
else:
arg_type = TypeConverter(self.errors,
line=line).visit(arg.annotation)
return Argument(Var(arg.arg), arg_type, self.visit(default), kind)
def add_method_to_class(
api: Union[SemanticAnalyzerPluginInterface, CheckerPluginInterface],
cls: ClassDef,
name: str,
args: List[Argument],
return_type: Type,
self_type: Optional[Type] = None,
tvar_def: Optional[TypeVarType] = None,
) -> None:
"""Adds a new method to a class definition."""
info = cls.info
# First remove any previously generated methods with the same name
# to avoid clashes and problems in the semantic analyzer.
if name in info.names:
sym = info.names[name]
if sym.plugin_generated and isinstance(sym.node, FuncDef):
cls.defs.body.remove(sym.node)
self_type = self_type or fill_typevars(info)
# TODO: semanal.py and checker.py seem to have subtly different implementations of
# named_type/named_generic_type (starting with the fact that we have to use different names
# for builtins), so it's easier to just check which one we're dealing with here and pick the
# correct function to use than to try to add a named_type method that behaves the same for
# both. We should probably combine those implementations at some point.
if isinstance(api, SemanticAnalyzerPluginInterface):
function_type = api.named_type('__builtins__.function')
else:
function_type = api.named_generic_type('builtins.function', [])
args = [Argument(Var('self'), self_type, None, ARG_POS)] + args
arg_types, arg_names, arg_kinds = [], [], []
for arg in args:
assert arg.type_annotation, 'All arguments must be fully typed.'
arg_types.append(arg.type_annotation)
arg_names.append(arg.variable.name)
arg_kinds.append(arg.kind)
signature = CallableType(arg_types, arg_kinds, arg_names, return_type,
function_type)
if tvar_def:
signature.variables = [tvar_def]
func = FuncDef(name, args, Block([PassStmt()]))
func.info = info
func.type = set_callable_name(signature, func)
func._fullname = info.fullname + '.' + name
func.line = info.line
# NOTE: we would like the plugin generated node to dominate, but we still
# need to keep any existing definitions so they get semantically analyzed.
if name in info.names:
# Get a nice unique name instead.
r_name = get_unique_redefinition_name(name, info.names)
info.names[r_name] = info.names[name]
info.names[name] = SymbolTableNode(MDEF, func, plugin_generated=True)
info.defn.defs.body.append(func)
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)