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 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 visit_mapping_pattern(self, o: MappingPattern) -> PatternType:
current_type = get_proper_type(self.type_context[-1])
can_match = True
captures: Dict[Expression, Type] = {}
for key, value in zip(o.keys, o.values):
inner_type = self.get_mapping_item_type(o, current_type, key)
if inner_type is None:
can_match = False
inner_type = self.chk.named_type("builtins.object")
pattern_type = self.accept(value, inner_type)
if is_uninhabited(pattern_type.type):
can_match = False
else:
self.update_type_map(captures, pattern_type.captures)
if o.rest is not None:
mapping = self.chk.named_type("typing.Mapping")
if is_subtype(current_type, mapping) and isinstance(current_type, Instance):
mapping_inst = map_instance_to_supertype(current_type, mapping.type)
dict_typeinfo = self.chk.lookup_typeinfo("builtins.dict")
dict_type = fill_typevars(dict_typeinfo)
rest_type = expand_type_by_instance(dict_type, mapping_inst)
else:
object_type = self.chk.named_type("builtins.object")
rest_type = self.chk.named_generic_type("builtins.dict",
[object_type, object_type])
captures[o.rest] = rest_type
if can_match:
# We can't narrow the type here, as Mapping key is invariant.
new_type = self.type_context[-1]
else:
new_type = UninhabitedType()
return PatternType(new_type, current_type, captures)
def map_type_from_supertype(typ: Type, sub_info: TypeInfo,
super_info: TypeInfo) -> Type:
"""Map type variables in a type defined in a supertype context to be valid
in the subtype context. Assume that the result is unique; if more than
one type is possible, return one of the alternatives.
For example, assume
. class D(Generic[S]) ...
. class C(D[E[T]], Generic[T]) ...
Now S in the context of D would be mapped to E[T] in the context of C.
"""
# Create the type of self in subtype, of form t[a1, ...].
inst_type = fill_typevars(sub_info)
if isinstance(inst_type, TupleType):
inst_type = tuple_fallback(inst_type)
# Map the type of self to supertype. This gets us a description of the
# supertype type variables in terms of subtype variables, i.e. t[t1, ...]
# so that any type variables in tN are to be interpreted in subtype
# context.
inst_type = map_instance_to_supertype(inst_type, super_info)
# Finally expand the type variables in type with those in the previously
# constructed type. Note that both type and inst_type may have type
# variables, but in type they are interpreted in supertype context while
# in inst_type they are interpreted in subtype context. This works even if
# the names of type variables in supertype and subtype overlap.
return expand_type_by_instance(typ, inst_type)
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 class_callable(init_type: CallableType,
info: TypeInfo,
type_type: Instance,
special_sig: Optional[str],
is_new: bool,
orig_self_type: Optional[Type] = None) -> CallableType:
"""Create a type object type based on the signature of __init__."""
variables: List[TypeVarLikeType] = []
variables.extend(info.defn.type_vars)
variables.extend(init_type.variables)
from mypy.subtypes import is_subtype
init_ret_type = get_proper_type(init_type.ret_type)
orig_self_type = get_proper_type(orig_self_type)
default_ret_type = fill_typevars(info)
explicit_type = init_ret_type if is_new else orig_self_type
if (isinstance(explicit_type, (Instance, TupleType))
# Only use the declared return type from __new__ or declared self in __init__
# if it is actually returning a subtype of what we would return otherwise.
and is_subtype(
explicit_type, default_ret_type, ignore_type_params=True)):
ret_type: Type = explicit_type
else:
ret_type = default_ret_type
callable_type = init_type.copy_modified(ret_type=ret_type,
fallback=type_type,
name=None,
variables=variables,
special_sig=special_sig)
c = callable_type.with_name(info.name)
return c
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 type_object_type_from_function(signature: FunctionLike, info: TypeInfo,
def_info: TypeInfo, fallback: Instance,
is_new: bool) -> FunctionLike:
# The __init__ method might come from a generic superclass
# (init_or_new.info) with type variables that do not map
# identically to the type variables of the class being constructed
# (info). For example
#
# class A(Generic[T]): def __init__(self, x: T) -> None: pass
# class B(A[List[T]], Generic[T]): pass
#
# We need to first map B's __init__ to the type (List[T]) -> None.
signature = bind_self(signature,
original_type=fill_typevars(info),
is_classmethod=is_new)
signature = cast(FunctionLike,
map_type_from_supertype(signature, info, def_info))
special_sig = None # type: Optional[str]
if def_info.fullname() == 'builtins.dict':
# Special signature!
special_sig = 'dict'
if isinstance(signature, CallableType):
return class_callable(signature, info, fallback, special_sig, is_new)
else:
# Overloaded __init__/__new__.
assert isinstance(signature, Overloaded)
items = [] # type: List[CallableType]
for item in signature.items():
items.append(
class_callable(item, info, fallback, special_sig, is_new))
return Overloaded(items)
def construct_sequence_child(self, outer_type: Type, inner_type: Type) -> Type:
"""
If outer_type is a child class of typing.Sequence returns a new instance of
outer_type, that is a Sequence of inner_type. If outer_type is not a child class of
typing.Sequence just returns a Sequence of inner_type
For example:
construct_sequence_child(List[int], str) = List[str]
"""
proper_type = get_proper_type(outer_type)
if isinstance(proper_type, UnionType):
types = [
self.construct_sequence_child(item, inner_type) for item in proper_type.items
if self.can_match_sequence(get_proper_type(item))
]
return make_simplified_union(types)
sequence = self.chk.named_generic_type("typing.Sequence", [inner_type])
if is_subtype(outer_type, self.chk.named_type("typing.Sequence")):
proper_type = get_proper_type(outer_type)
assert isinstance(proper_type, Instance)
empty_type = fill_typevars(proper_type.type)
partial_type = expand_type_by_instance(empty_type, sequence)
return expand_type_by_instance(partial_type, proper_type)
else:
return sequence
def callable_type(fdef: FuncItem,
fallback: Instance,
ret_type: Optional[Type] = None) -> CallableType:
# TODO: somewhat unfortunate duplication with prepare_method_signature in semanal
if fdef.info and not fdef.is_static and fdef.arg_names:
self_type: Type = fill_typevars(fdef.info)
if fdef.is_class or fdef.name == '__new__':
self_type = TypeType.make_normalized(self_type)
args = [
self_type
] + [AnyType(TypeOfAny.unannotated)] * (len(fdef.arg_names) - 1)
else:
args = [AnyType(TypeOfAny.unannotated)] * len(fdef.arg_names)
return CallableType(
args,
fdef.arg_kinds,
fdef.arg_names,
ret_type or AnyType(TypeOfAny.unannotated),
fallback,
name=fdef.name,
line=fdef.line,
column=fdef.column,
implicit=True,
)
def map_type_from_supertype(typ: Type,
sub_info: TypeInfo,
super_info: TypeInfo) -> Type:
"""Map type variables in a type defined in a supertype context to be valid
in the subtype context. Assume that the result is unique; if more than
one type is possible, return one of the alternatives.
For example, assume
. class D(Generic[S]) ...
. class C(D[E[T]], Generic[T]) ...
Now S in the context of D would be mapped to E[T] in the context of C.
"""
# Create the type of self in subtype, of form t[a1, ...].
inst_type = fill_typevars(sub_info)
if isinstance(inst_type, TupleType):
inst_type = tuple_fallback(inst_type)
# Map the type of self to supertype. This gets us a description of the
# supertype type variables in terms of subtype variables, i.e. t[t1, ...]
# so that any type variables in tN are to be interpreted in subtype
# context.
inst_type = map_instance_to_supertype(inst_type, super_info)
# Finally expand the type variables in type with those in the previously
# constructed type. Note that both type and inst_type may have type
# variables, but in type they are interpreted in supertype context while
# in inst_type they are interpreted in subtype context. This works even if
# the names of type variables in supertype and subtype overlap.
return expand_type_by_instance(typ, inst_type)
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 supported_self_type(typ: ProperType) -> bool:
"""Is this a supported kind of explicit self-types?
Currently, this means a X or Type[X], where X is an instance or
a type variable with an instance upper bound.
"""
if isinstance(typ, TypeType):
return supported_self_type(typ.item)
return (isinstance(typ, TypeVarType)
or (isinstance(typ, Instance) and typ != fill_typevars(typ.type)))
def class_callable(init_type: CallableType, info: TypeInfo, type_type: Instance,
special_sig: Optional[str]) -> CallableType:
"""Create a type object type based on the signature of __init__."""
variables = [] # type: List[TypeVarDef]
variables.extend(info.defn.type_vars)
variables.extend(init_type.variables)
callable_type = init_type.copy_modified(
ret_type=fill_typevars(info), fallback=type_type, name=None, variables=variables,
special_sig=special_sig)
c = callable_type.with_name(info.name())
return c
def class_callable(init_type: CallableType, info: TypeInfo, type_type: Instance,
special_sig: Optional[str]) -> CallableType:
"""Create a type object type based on the signature of __init__."""
variables = [] # type: List[TypeVarDef]
variables.extend(info.defn.type_vars)
variables.extend(init_type.variables)
callable_type = init_type.copy_modified(
ret_type=fill_typevars(info), fallback=type_type, name=None, variables=variables,
special_sig=special_sig)
c = callable_type.with_name(info.name())
return c
def _transform_class(context: ClassDefContext) -> None:
cls = context.cls
instanceType = fill_typevars(cls.info)
assert isinstance(instanceType, mypy.types.Instance)
cases = _get_and_delete_cases(context)
for case in cases:
_add_constructor_for_case(context, case, selfType=instanceType)
_add_accessor_for_case(context, case)
_add_match(context, cases)
def add_method_to_class(
api: SemanticAnalyzerPluginInterface,
cls: ClassDef,
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 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)
function_type = 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
# 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 type_object_type_from_function(signature: FunctionLike, info: TypeInfo,
def_info: TypeInfo, fallback: Instance,
is_new: bool) -> FunctionLike:
# We first need to record all non-trivial (explicit) self types in __init__,
# since they will not be available after we bind them. Note, we use explicit
# self-types only in the defining class, similar to __new__ (but not exactly the same,
# see comment in class_callable below). This is mostly useful for annotating library
# classes such as subprocess.Popen.
default_self = fill_typevars(info)
if not is_new and not info.is_newtype:
orig_self_types = [(it.arg_types[0]
if it.arg_types and it.arg_types[0] != default_self
and it.arg_kinds[0] == ARG_POS else None)
for it in signature.items]
else:
orig_self_types = [None] * len(signature.items)
# The __init__ method might come from a generic superclass 'def_info'
# with type variables that do not map identically to the type variables of
# the class 'info' being constructed. For example:
#
# class A(Generic[T]):
# def __init__(self, x: T) -> None: ...
# class B(A[List[T]]):
# ...
#
# We need to map B's __init__ to the type (List[T]) -> None.
signature = bind_self(signature,
original_type=default_self,
is_classmethod=is_new)
signature = cast(FunctionLike,
map_type_from_supertype(signature, info, def_info))
special_sig: Optional[str] = None
if def_info.fullname == 'builtins.dict':
# Special signature!
special_sig = 'dict'
if isinstance(signature, CallableType):
return class_callable(signature, info, fallback, special_sig, is_new,
orig_self_types[0])
else:
# Overloaded __init__/__new__.
assert isinstance(signature, Overloaded)
items: List[CallableType] = []
for item, orig_self in zip(signature.items, orig_self_types):
items.append(
class_callable(item, info, fallback, special_sig, is_new,
orig_self))
return Overloaded(items)
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 _transform_class(context: ClassDefContext) -> None:
cls = context.cls
instanceType = fill_typevars(cls.info)
assert isinstance(instanceType, mypy.types.Instance)
cases = _get_and_delete_cases(context)
if cases is None: # Cases were not successfully deleted. We need to defer
context.api.defer()
return
for case in cases:
_add_constructor_for_case(context, case, selfType=instanceType)
_add_accessor_for_case(context, case)
_add_match(context, cases)
def construct_sequence_child(self, outer_type: Type,
inner_type: Type) -> Type:
"""
If outer_type is a child class of typing.Sequence returns a new instance of
outer_type, that is a Sequence of inner_type. If outer_type is not a child class of
typing.Sequence just returns a Sequence of inner_type
For example:
construct_sequence_child(List[int], str) = List[str]
"""
sequence = self.chk.named_generic_type("typing.Sequence", [inner_type])
if is_subtype(outer_type, self.chk.named_type("typing.Sequence")):
proper_type = get_proper_type(outer_type)
assert isinstance(proper_type, Instance)
empty_type = fill_typevars(proper_type.type)
partial_type = expand_type_by_instance(empty_type, sequence)
return expand_type_by_instance(partial_type, proper_type)
else:
return sequence
def _lookup_type(self, fullname: str) -> types.Type:
ctx = self._ctx
type_sym = ctx.api.lookup_fully_qualified_or_none(fullname)
if type_sym is None:
raise DeferException
t: types.Type
if isinstance(type_sym.node, nodes.TypeInfo):
from mypy.typevars import fill_typevars
t = fill_typevars(type_sym.node)
elif type_sym.type:
t = type_sym.type
else:
ctx.api.fail(f'cannot find {fullname}', ctx.cls)
return t
def create_ortho_diff_class(base1: TypeInfo, base2: TypeInfo,
api: SemanticAnalyzerPluginInterface,
call_ctx: Context) -> Tuple[str, SymbolTableNode]:
# https://github.com/dropbox/sqlalchemy-stubs/blob/55470ceab8149db983411d5c094c9fe16343c58b/sqlmypy.py#L173-L216
cls_name = get_ortho_diff_name(base1.defn, base2.defn)
class_def = ClassDef(cls_name, Block([]))
class_def.fullname = api.qualified_name(cls_name)
info = TypeInfo(SymbolTable(), class_def, api.cur_mod_id)
class_def.info = info
obj = api.builtin_type('builtins.object')
info.bases = [cast(Instance, fill_typevars(b)) for b in (base1, base2)]
try:
calculate_mro(info)
except MroError:
api.fail('Unable to calculate MRO for dynamic class', call_ctx)
info.bases = [obj]
info.fallback_to_any = True
return cls_name, SymbolTableNode(GDEF, info)
def class_callable(init_type: CallableType, info: TypeInfo,
type_type: Instance, special_sig: Optional[str],
is_new: bool) -> CallableType:
"""Create a type object type based on the signature of __init__."""
variables = [] # type: List[TypeVarDef]
variables.extend(info.defn.type_vars)
variables.extend(init_type.variables)
init_ret_type = get_proper_type(init_type.ret_type)
if is_new and isinstance(init_ret_type, (Instance, TupleType)):
ret_type = init_type.ret_type # type: Type
else:
ret_type = fill_typevars(info)
callable_type = init_type.copy_modified(ret_type=ret_type,
fallback=type_type,
name=None,
variables=variables,
special_sig=special_sig)
c = callable_type.with_name(info.name())
return c
def add_method(
ctx: ClassDefContext,
name: str,
args: List[Argument],
return_type: Type,
self_type: Optional[Type] = None,
tvar_def: Optional[TypeVarDef] = None,
is_classmethod: bool = False,
is_new: bool = False,
# is_staticmethod: bool = False,
) -> None:
"""
Adds a new method to a class.
This can be dropped if/when https://github.com/python/mypy/issues/7301 is merged
"""
info = ctx.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):
ctx.cls.defs.body.remove(sym.node)
self_type = self_type or fill_typevars(info)
if is_classmethod or is_new:
first = [
Argument(Var("_cls"), TypeType.make_normalized(self_type), None,
ARG_POS)
]
# elif is_staticmethod:
# first = []
else:
self_type = self_type or fill_typevars(info)
first = [Argument(Var("self"), self_type, None, ARG_POS)]
args = first + 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(get_name(arg.variable))
arg_kinds.append(arg.kind)
function_type = ctx.api.named_type("__builtins__.function")
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.is_class = is_classmethod
# func.is_static = is_staticmethod
func._fullname = get_fullname(info) + "." + 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]
if is_classmethod: # or is_staticmethod:
func.is_decorated = True
v = Var(name, func.type)
v.info = info
v._fullname = func._fullname
# if is_classmethod:
v.is_classmethod = True
dec = Decorator(func, [NameExpr("classmethod")], v)
# else:
# v.is_staticmethod = True
# dec = Decorator(func, [NameExpr('staticmethod')], v)
dec.line = info.line
sym = SymbolTableNode(MDEF, dec)
else:
sym = SymbolTableNode(MDEF, func)
sym.plugin_generated = True
info.names[name] = sym
info.defn.defs.body.append(func)
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")
bool_type = cls_def_ctx.api.named_type("__builtins__.bool")
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, # type: ignore # i don't understand this one but it works
[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
cat_return_type = fill_typevars(info)
maybe_cat_return_type = make_optional(cat_return_type)
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,
)
],
cat_return_type,
)
if TRY_STRUCTURE_NAME not in info.names:
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,
)
],
maybe_cat_return_type,
)
if UNSTRUCTURE_NAME not in info.names:
add_method(
cls_def_ctx,
UNSTRUCTURE_NAME,
[
Argument(
Var("strip_defaults", bool_type),
bool_type,
None,
ARG_NAMED_OPT,
)
],
dict_type,
)
def _add_method(
ctx: ClassDefContext,
name: str,
args: List[Argument],
return_type: mypy.types.Type,
self_type: Optional[mypy.types.Type] = None,
tvar_def: Optional[mypy.types.TypeVarDef] = None,
is_classmethod: bool = False,
) -> None:
"""Adds a new method to a class.
"""
info = ctx.cls.info
# First remove any previously generated methods with the same name
# to avoid clashes and problems in new semantic analyzer.
if name in info.names:
sym = info.names[name]
if sym.plugin_generated and isinstance(sym.node, FuncDef):
ctx.cls.defs.body.remove(sym.node)
if is_classmethod:
first = Argument(
Var('cls'),
# Working around python/mypy#5416.
# This should be: mypy.types.TypeType.make_normalized(self_type)
mypy.types.AnyType(mypy.types.TypeOfAny.implementation_artifact),
None,
ARG_POS)
else:
self_type = self_type or fill_typevars(info)
first = Argument(Var('self'), self_type, None, ARG_POS)
args = [first] + args
function_type = ctx.api.named_type('__builtins__.function')
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(get_name(arg.variable))
arg_kinds.append(arg.kind)
signature = mypy.types.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.is_class = is_classmethod
func.type = set_callable_name(signature, func)
func._fullname = get_fullname(info) + '.' + 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.defn.defs.body.append(func)
info.names[name] = SymbolTableNode(MDEF, func, plugin_generated=True)
def __init__(self, ctx: 'mypy.plugin.ClassDefContext') -> None:
self.ctx = ctx
self.self_type = fill_typevars(ctx.cls.info)
def __init__(self, info: TypeInfo, function_type: Instance) -> None:
self.info = info
self.self_type = fill_typevars(info)
self.function_type = function_type
def update_dataclass_json_type(ctx: ClassDefContext) -> None:
str_type = builtin_type(ctx.api, "str")
int_type = builtin_type(ctx.api, "int")
sep_type = UnionType.make_union(
[builtin_type(ctx.api, 'tuple', [str_type, str_type]),
NoneType()])
indent_type = UnionType.make_union([int_type, NoneType()])
args = [
Argument(Var('separators', sep_type), sep_type, None, ARG_NAMED_OPT),
Argument(Var('indent', indent_type), indent_type, None, ARG_NAMED_OPT),
]
add_method_to_class(ctx.api,
ctx.cls,
'to_json',
args=args,
return_type=str_type)
# It would be lovely to actually have this return a typed dict ;)
json_dict_type = builtin_type(
ctx.api, 'dict', [str_type, AnyType(TypeOfAny.explicit)])
add_method_to_class(ctx.api,
ctx.cls,
'to_dict',
args=[],
return_type=json_dict_type)
instance_type = fill_typevars(ctx.cls.info)
bool_type = builtin_type(ctx.api, 'bool')
args = [
Argument(Var('o', json_dict_type), json_dict_type, None, ARG_POS),
Argument(Var('infer_missing', bool_type), bool_type, None,
ARG_NAMED_OPT),
]
add_classmethod_to_class(
ctx.api,
ctx.cls,
'from_dict',
args=args,
return_type=instance_type,
)
bytes_type = builtin_type(ctx.api, 'bytes')
json_data_type = UnionType.make_union([str_type, bytes_type])
args = [
Argument(Var('json_data', json_data_type), json_data_type, None,
ARG_POS),
Argument(Var('infer_missing', bool_type), bool_type, None,
ARG_NAMED_OPT),
]
add_classmethod_to_class(
ctx.api,
ctx.cls,
'from_json',
args=args,
return_type=instance_type,
)
def __init__(self, ctx: 'mypy.plugin.ClassDefContext') -> None:
self.ctx = ctx
self.self_type = fill_typevars(ctx.cls.info)
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 __init__(self, info: TypeInfo, function_type: Instance) -> None:
self.info = info
self.self_type = fill_typevars(info)
self.function_type = function_type
def add_method_to_class(
api: SemanticAnalyzerPluginInterface,
cls: ClassDef,
name: str,
args: List[Argument],
return_type: Type,
self_type: Optional[Type] = None,
tvar_def: Optional[TypeVarDef] = None,
is_classmethod: bool = False,
) -> None:
"""
Adds a new method to a class definition.
NOTE:
Copied from mypy/plugins/common.py and extended with support for adding
classmethods based on https://github.com/python/mypy/pull/7796
"""
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)
# Add either self or cls as the first argument
if is_classmethod:
first = Argument(Var("cls"), TypeType.make_normalized(self_type), None,
ARG_POS)
else:
first = Argument(Var("self"), self_type, None, ARG_POS)
args = [first] + 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)
function_type = api.named_type("__builtins__.function")
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 # pylint: disable=protected-access
func.line = info.line
func.is_class = is_classmethod
# 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]
if is_classmethod:
func.is_decorated = True
v = Var(name, func.type)
v.info = info
v._fullname = func._fullname # pylint: disable=protected-access
v.is_classmethod = True
dec = Decorator(func, [NameExpr("classmethod")], v)
dec.line = info.line
sym = SymbolTableNode(MDEF, dec)
else:
sym = SymbolTableNode(MDEF, func)
sym.plugin_generated = True
info.names[name] = sym
info.defn.defs.body.append(func)
