mirror of
https://github.com/davidhalter/django-stubs.git
synced 2025-12-07 12:44:29 +08:00
232 lines
9.4 KiB
Python
232 lines
9.4 KiB
Python
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'
|