from typing import Optional, cast from mypy.checker import TypeChecker from mypy.nodes import ListExpr, NameExpr, StrExpr, TupleExpr, TypeInfo, Expression from mypy.plugin import FunctionContext from mypy.types import ( AnyType, CallableType, Instance, TupleType, Type, UnionType, ) from mypy_django_plugin.lib import fullnames, helpers, metadata def extract_referred_to_type(ctx: FunctionContext) -> Optional[Instance]: api = cast(TypeChecker, ctx.api) if 'to' not in ctx.callee_arg_names: api.msg.fail(f'to= parameter must be set for {ctx.context.callee.fullname}', context=ctx.context) return None arg_type = ctx.arg_types[ctx.callee_arg_names.index('to')][0] if not isinstance(arg_type, CallableType): to_arg_expr = ctx.args[ctx.callee_arg_names.index('to')][0] if not isinstance(to_arg_expr, StrExpr): # not string, not supported return None try: model_fullname = helpers.get_model_fullname_from_string(to_arg_expr.value, all_modules=api.modules) except helpers.SelfReference: model_fullname = api.tscope.classes[-1].fullname() except helpers.SameFileModel as exc: model_fullname = api.tscope.classes[-1].module_name + '.' + exc.model_cls_name if model_fullname is None: return None model_info = helpers.lookup_fully_qualified_generic(model_fullname, all_modules=api.modules) if model_info is None or not isinstance(model_info, TypeInfo): return None return Instance(model_info, []) referred_to_type = arg_type.ret_type if not isinstance(referred_to_type, Instance): return None if not referred_to_type.type.has_base(fullnames.MODEL_CLASS_FULLNAME): ctx.api.msg.fail(f'to= parameter value must be ' f'a subclass of {fullnames.MODEL_CLASS_FULLNAME!r}', context=ctx.context) return None return referred_to_type def convert_any_to_type(typ: Type, referred_to_type: Type) -> Type: if isinstance(typ, UnionType): converted_items = [] for item in typ.items: converted_items.append(convert_any_to_type(item, referred_to_type)) return UnionType.make_union(converted_items, line=typ.line, column=typ.column) if isinstance(typ, Instance): args = [] for default_arg in typ.args: if isinstance(default_arg, AnyType): args.append(referred_to_type) else: args.append(default_arg) return helpers.reparametrize_instance(typ, args) if isinstance(typ, AnyType): return referred_to_type return typ def fill_descriptor_types_for_related_field(ctx: FunctionContext) -> Type: default_return_type = set_descriptor_types_for_field(ctx) referred_to_type = extract_referred_to_type(ctx) if referred_to_type is None: return default_return_type # replace Any with referred_to_type args = [] for default_arg in default_return_type.args: args.append(convert_any_to_type(default_arg, referred_to_type)) return helpers.reparametrize_instance(ctx.default_return_type, new_args=args) def set_descriptor_types_for_field(ctx: FunctionContext) -> Instance: default_return_type = cast(Instance, ctx.default_return_type) is_nullable = helpers.parse_bool(helpers.get_call_argument_by_name(ctx, 'null')) set_type = helpers.get_private_descriptor_type(default_return_type.type, '_pyi_private_set_type', is_nullable=is_nullable) get_type = helpers.get_private_descriptor_type(default_return_type.type, '_pyi_private_get_type', is_nullable=is_nullable) return helpers.reparametrize_instance(default_return_type, [set_type, get_type]) def determine_type_of_array_field(ctx: FunctionContext) -> Type: default_return_type = set_descriptor_types_for_field(ctx) base_field_arg_type = helpers.get_call_argument_type_by_name(ctx, 'base_field') if not base_field_arg_type or not isinstance(base_field_arg_type, Instance): return default_return_type base_type = base_field_arg_type.args[1] # extract __get__ type args = [] for default_arg in default_return_type.args: args.append(convert_any_to_type(default_arg, base_type)) return helpers.reparametrize_instance(default_return_type, args) def transform_into_proper_return_type(ctx: FunctionContext) -> Type: default_return_type = ctx.default_return_type if not isinstance(default_return_type, Instance): return default_return_type if helpers.has_any_of_bases(default_return_type.type, fullnames.RELATED_FIELDS_CLASSES): return fill_descriptor_types_for_related_field(ctx) if default_return_type.type.has_base(fullnames.ARRAY_FIELD_FULLNAME): return determine_type_of_array_field(ctx) return set_descriptor_types_for_field(ctx) def process_field_instantiation(ctx: FunctionContext) -> Type: # Parse __init__ parameters of field into corresponding Model's metadata parse_field_init_arguments_into_model_metadata(ctx) return transform_into_proper_return_type(ctx) def _parse_choices_type(ctx: FunctionContext, choices_arg: Expression) -> Optional[str]: if isinstance(choices_arg, (TupleExpr, ListExpr)): # iterable of 2 element tuples of two kinds _, analyzed_choices = ctx.api.analyze_iterable_item_type(choices_arg) if isinstance(analyzed_choices, TupleType): first_element_type = analyzed_choices.items[0] if isinstance(first_element_type, Instance): return first_element_type.type.fullname() def _parse_referenced_model(ctx: FunctionContext, to_arg: Expression) -> Optional[TypeInfo]: if isinstance(to_arg, NameExpr) and isinstance(to_arg.node, TypeInfo): # reference to the model class return to_arg.node elif isinstance(to_arg, StrExpr): referenced_model_info = helpers.get_model_info(to_arg.value, ctx.api.modules) if referenced_model_info is not None: return referenced_model_info def parse_field_init_arguments_into_model_metadata(ctx: FunctionContext) -> None: outer_model = ctx.api.scope.active_class() if outer_model is None or not outer_model.has_base(fullnames.MODEL_CLASS_FULLNAME): # outside models.Model class, undetermined return # Determine name of the current field for attr_name, stmt in helpers.iter_over_class_level_assignments(outer_model.defn): if stmt == ctx.context: field_name = attr_name break else: return model_fields_metadata = metadata.get_fields_metadata(outer_model) # primary key is_primary_key = False primary_key_arg = helpers.get_call_argument_by_name(ctx, 'primary_key') if primary_key_arg: is_primary_key = helpers.parse_bool(primary_key_arg) model_fields_metadata[field_name] = {'primary_key': is_primary_key} # choices choices_arg = helpers.get_call_argument_by_name(ctx, 'choices') if choices_arg: choices_type_fullname = _parse_choices_type(ctx.api, choices_arg) if choices_type_fullname: model_fields_metadata[field_name]['choices_type'] = choices_type_fullname # nullability null_arg = helpers.get_call_argument_by_name(ctx, 'null') is_nullable = False if null_arg: is_nullable = helpers.parse_bool(null_arg) model_fields_metadata[field_name]['null'] = is_nullable # is_blankable blank_arg = helpers.get_call_argument_by_name(ctx, 'blank') is_blankable = False if blank_arg: is_blankable = helpers.parse_bool(blank_arg) model_fields_metadata[field_name]['blank'] = is_blankable # default default_arg = helpers.get_call_argument_by_name(ctx, 'default') if default_arg and not helpers.is_none_expr(default_arg): model_fields_metadata[field_name]['default_specified'] = True if helpers.has_any_of_bases(ctx.default_return_type.type, fullnames.RELATED_FIELDS_CLASSES): # to to_arg = helpers.get_call_argument_by_name(ctx, 'to') if to_arg: referenced_model = _parse_referenced_model(ctx, to_arg) if referenced_model is not None: model_fields_metadata[field_name]['to'] = referenced_model.fullname() else: model_fields_metadata[field_name]['to'] = to_arg.value # referenced_model = to_arg.value # raise helpers.IncompleteDefnException() # model_fields_metadata[field_name]['to'] = referenced_model.fullname() # if referenced_model is not None: # model_fields_metadata[field_name]['to'] = referenced_model.fullname() # else: # assert isinstance(to_arg, StrExpr) # model_fields_metadata[field_name]['to'] = to_arg.value # related_name related_name_arg = helpers.get_call_argument_by_name(ctx, 'related_name') if related_name_arg: if isinstance(related_name_arg, StrExpr): model_fields_metadata[field_name]['related_name'] = related_name_arg.value else: model_fields_metadata[field_name]['related_name'] = outer_model.name().lower() + '_set'