def extract_proper_type_queryset_values(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
# called on QuerySet, return QuerySet of something
assert isinstance(ctx.type, Instance)
assert isinstance(ctx.default_return_type, Instance)
model_type = _extract_model_type_from_queryset(ctx.type)
if model_type is None:
return AnyType(TypeOfAny.from_omitted_generics)
model_cls = django_context.get_model_class_by_fullname(model_type.type.fullname)
if model_cls is None:
return ctx.default_return_type
field_lookups = resolve_field_lookups(ctx.args[0], django_context)
if field_lookups is None:
return AnyType(TypeOfAny.from_error)
if len(field_lookups) == 0:
for field in django_context.get_model_fields(model_cls):
field_lookups.append(field.attname)
column_types: 'OrderedDict[str, MypyType]' = OrderedDict()
for field_lookup in field_lookups:
field_lookup_type = get_field_type_from_lookup(ctx, django_context, model_cls,
lookup=field_lookup, method='values')
if field_lookup_type is None:
return helpers.reparametrize_instance(ctx.default_return_type, [model_type, AnyType(TypeOfAny.from_error)])
column_types[field_lookup] = field_lookup_type
row_type = helpers.make_typeddict(ctx.api, column_types, set(column_types.keys()))
return helpers.reparametrize_instance(ctx.default_return_type, [model_type, row_type])
def get_field_type_from_lookup(
ctx: MethodContext,
django_context: DjangoContext,
model_cls: Type[Model],
*,
method: str,
lookup: str,
silent_on_error: bool = False,
) -> Optional[MypyType]:
try:
lookup_field = django_context.resolve_lookup_into_field(
model_cls, lookup)
except FieldError as exc:
if not silent_on_error:
ctx.api.fail(exc.args[0], ctx.context)
return None
except LookupsAreUnsupported:
return AnyType(TypeOfAny.explicit)
if (isinstance(lookup_field, RelatedField)
and lookup_field.column == lookup) or isinstance(
lookup_field, ForeignObjectRel):
related_model_cls = django_context.get_field_related_model_cls(
lookup_field)
if related_model_cls is None:
return AnyType(TypeOfAny.from_error)
lookup_field = django_context.get_primary_key_field(related_model_cls)
field_get_type = django_context.get_field_get_type(
helpers.get_typechecker_api(ctx), lookup_field, method=method)
return field_get_type
def __init__(self, options: Options) -> None:
super().__init__(options)
django_settings_module = extract_django_settings_module(options.config_file)
# Add paths from MYPYPATH env var
sys.path.extend(mypy_path())
# Add paths from mypy_path config option
sys.path.extend(options.mypy_path)
self.django_context = DjangoContext(django_settings_module)
def get_values_list_row_type(ctx: MethodContext, django_context: DjangoContext, model_cls: Type[Model],
flat: bool, named: bool) -> MypyType:
field_lookups = resolve_field_lookups(ctx.args[0], django_context)
if field_lookups is None:
return AnyType(TypeOfAny.from_error)
typechecker_api = helpers.get_typechecker_api(ctx)
if len(field_lookups) == 0:
if flat:
primary_key_field = django_context.get_primary_key_field(model_cls)
lookup_type = get_field_type_from_lookup(ctx, django_context, model_cls,
lookup=primary_key_field.attname, method='values_list')
assert lookup_type is not None
return lookup_type
elif named:
column_types: 'OrderedDict[str, MypyType]' = OrderedDict()
for field in django_context.get_model_fields(model_cls):
column_type = django_context.get_field_get_type(typechecker_api, field,
method='values_list')
column_types[field.attname] = column_type
return helpers.make_oneoff_named_tuple(typechecker_api, 'Row', column_types)
else:
# flat=False, named=False, all fields
field_lookups = []
for field in django_context.get_model_fields(model_cls):
field_lookups.append(field.attname)
if len(field_lookups) > 1 and flat:
typechecker_api.fail("'flat' is not valid when 'values_list' is called with more than one field", ctx.context)
return AnyType(TypeOfAny.from_error)
column_types = OrderedDict()
for field_lookup in field_lookups:
lookup_field_type = get_field_type_from_lookup(ctx, django_context, model_cls,
lookup=field_lookup, method='values_list')
if lookup_field_type is None:
return AnyType(TypeOfAny.from_error)
column_types[field_lookup] = lookup_field_type
if flat:
assert len(column_types) == 1
row_type = next(iter(column_types.values()))
elif named:
row_type = helpers.make_oneoff_named_tuple(typechecker_api, 'Row', column_types)
else:
row_type = helpers.make_tuple(typechecker_api, list(column_types.values()))
return row_type
def return_proper_field_type_from_get_field(
ctx: MethodContext, django_context: DjangoContext) -> MypyType:
# Options instance
assert isinstance(ctx.type, Instance)
model_type = ctx.type.args[0]
if not isinstance(model_type, Instance):
return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
model_cls = django_context.get_model_class_by_fullname(
model_type.type.fullname())
if model_cls is None:
return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
field_name_expr = helpers.get_call_argument_by_name(ctx, 'field_name')
if field_name_expr is None:
return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
field_name = helpers.resolve_string_attribute_value(
field_name_expr, ctx, django_context)
if field_name is None:
return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
try:
field = model_cls._meta.get_field(field_name)
except FieldDoesNotExist as exc:
ctx.api.fail(exc.args[0], ctx.context)
return AnyType(TypeOfAny.from_error)
field_fullname = helpers.get_class_fullname(field.__class__)
return _get_field_instance(ctx, field_fullname)
def _get_current_field_from_assignment(
ctx: FunctionContext,
django_context: DjangoContext) -> Optional[Field]:
outer_model_info = helpers.get_typechecker_api(ctx).scope.active_class()
if (outer_model_info is None or not helpers.is_model_subclass_info(
outer_model_info, django_context)):
return None
field_name = None
for stmt in outer_model_info.defn.defs.body:
if isinstance(stmt, AssignmentStmt):
if stmt.rvalue == ctx.context:
if not isinstance(stmt.lvalues[0], NameExpr):
return None
field_name = stmt.lvalues[0].name
break
if field_name is None:
return None
model_cls = django_context.get_model_class_by_fullname(
outer_model_info.fullname)
if model_cls is None:
return None
current_field = model_cls._meta.get_field(field_name)
return current_field
def resolve_combinable_type(combinable_type: Instance,
django_context: DjangoContext) -> MypyType:
if combinable_type.type.fullname != fullnames.F_EXPRESSION_FULLNAME:
# Combinables aside from F expressions are unsupported
return AnyType(TypeOfAny.explicit)
return django_context.resolve_f_expression_type(combinable_type)
def typecheck_queryset_filter(ctx: MethodContext,
django_context: DjangoContext) -> MypyType:
# Expected formal arguments for filter methods are `*args` and `**kwargs`. We'll only typecheck
# `**kwargs`, which means that `arg_names[1]` is what we're interested in.
lookup_kwargs = ctx.arg_names[1]
provided_lookup_types = ctx.arg_types[1]
assert isinstance(ctx.type, Instance)
if not ctx.type.args or not isinstance(ctx.type.args[0], Instance):
return ctx.default_return_type
model_cls_fullname = ctx.type.args[0].type.fullname
model_cls = django_context.get_model_class_by_fullname(model_cls_fullname)
if model_cls is None:
return ctx.default_return_type
for lookup_kwarg, provided_type in zip(lookup_kwargs,
provided_lookup_types):
if lookup_kwarg is None:
continue
if isinstance(provided_type, Instance) and provided_type.type.has_base(
fullnames.COMBINABLE_EXPRESSION_FULLNAME):
provided_type = resolve_combinable_type(provided_type,
django_context)
lookup_type: MypyType
if is_annotated_model_fullname(model_cls_fullname):
lookup_type = AnyType(TypeOfAny.implementation_artifact)
else:
lookup_type = django_context.resolve_lookup_expected_type(
ctx, model_cls, lookup_kwarg)
# Managers as provided_type is not supported yet
if isinstance(provided_type, Instance) and helpers.has_any_of_bases(
provided_type.type, (fullnames.MANAGER_CLASS_FULLNAME,
fullnames.QUERYSET_CLASS_FULLNAME)):
return ctx.default_return_type
helpers.check_types_compatible(
ctx,
expected_type=lookup_type,
actual_type=provided_type,
error_message=f"Incompatible type for lookup {lookup_kwarg!r}:",
)
return ctx.default_return_type
def typecheck_queryset_filter(ctx: MethodContext,
django_context: DjangoContext) -> MypyType:
lookup_kwargs = ctx.arg_names[1]
provided_lookup_types = ctx.arg_types[1]
assert isinstance(ctx.type, Instance)
if not ctx.type.args or not isinstance(ctx.type.args[0], Instance):
return ctx.default_return_type
model_cls_fullname = ctx.type.args[0].type.fullname
model_cls = django_context.get_model_class_by_fullname(model_cls_fullname)
if model_cls is None:
return ctx.default_return_type
for lookup_kwarg, provided_type in zip(lookup_kwargs,
provided_lookup_types):
if lookup_kwarg is None:
continue
if isinstance(provided_type, Instance) and provided_type.type.has_base(
"django.db.models.expressions.Combinable"):
provided_type = resolve_combinable_type(provided_type,
django_context)
lookup_type: MypyType
if is_annotated_model_fullname(model_cls_fullname):
lookup_type = AnyType(TypeOfAny.implementation_artifact)
else:
lookup_type = django_context.resolve_lookup_expected_type(
ctx, model_cls, lookup_kwarg)
# Managers as provided_type is not supported yet
if isinstance(provided_type, Instance) and helpers.has_any_of_bases(
provided_type.type, (fullnames.MANAGER_CLASS_FULLNAME,
fullnames.QUERYSET_CLASS_FULLNAME)):
return ctx.default_return_type
helpers.check_types_compatible(
ctx,
expected_type=lookup_type,
actual_type=provided_type,
error_message=f"Incompatible type for lookup {lookup_kwarg!r}:",
)
return ctx.default_return_type
def redefine_and_typecheck_model_init(
ctx: FunctionContext, django_context: DjangoContext) -> MypyType:
assert isinstance(ctx.default_return_type, Instance)
model_fullname = ctx.default_return_type.type.fullname
model_cls = django_context.get_model_class_by_fullname(model_fullname)
if model_cls is None:
return ctx.default_return_type
return typecheck_model_method(ctx, django_context, model_cls, '__init__')
def fill_descriptor_types_for_related_field(ctx: FunctionContext, django_context: DjangoContext) -> MypyType:
current_field = _get_current_field_from_assignment(ctx, django_context)
if current_field is None:
return AnyType(TypeOfAny.from_error)
assert isinstance(current_field, RelatedField)
related_model_cls = django_context.get_field_related_model_cls(current_field)
if related_model_cls is None:
return AnyType(TypeOfAny.from_error)
default_related_field_type = set_descriptor_types_for_field(ctx)
# self reference with abstract=True on the model where ForeignKey is defined
current_model_cls = current_field.model
if (current_model_cls._meta.abstract
and current_model_cls == related_model_cls):
# for all derived non-abstract classes, set variable with this name to
# __get__/__set__ of ForeignKey of derived model
for model_cls in django_context.all_registered_model_classes:
if issubclass(model_cls, current_model_cls) and not model_cls._meta.abstract:
derived_model_info = helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx), model_cls)
if derived_model_info is not None:
fk_ref_type = Instance(derived_model_info, [])
derived_fk_type = reparametrize_related_field_type(default_related_field_type,
set_type=fk_ref_type, get_type=fk_ref_type)
helpers.add_new_sym_for_info(derived_model_info,
name=current_field.name,
sym_type=derived_fk_type)
related_model = related_model_cls
related_model_to_set = related_model_cls
if related_model_to_set._meta.proxy_for_model is not None:
related_model_to_set = related_model_to_set._meta.proxy_for_model
typechecker_api = helpers.get_typechecker_api(ctx)
related_model_info = helpers.lookup_class_typeinfo(typechecker_api, related_model)
if related_model_info is None:
# maybe no type stub
related_model_type = AnyType(TypeOfAny.unannotated)
else:
related_model_type = Instance(related_model_info, []) # type: ignore
related_model_to_set_info = helpers.lookup_class_typeinfo(typechecker_api, related_model_to_set)
if related_model_to_set_info is None:
# maybe no type stub
related_model_to_set_type = AnyType(TypeOfAny.unannotated)
else:
related_model_to_set_type = Instance(related_model_to_set_info, []) # type: ignore
# replace Any with referred_to_type
return reparametrize_related_field_type(default_related_field_type,
set_type=related_model_to_set_type,
get_type=related_model_type)
def redefine_and_typecheck_model_create(
ctx: MethodContext, django_context: DjangoContext) -> MypyType:
if not isinstance(ctx.default_return_type, Instance):
# only work with ctx.default_return_type = model Instance
return ctx.default_return_type
model_fullname = ctx.default_return_type.type.fullname
model_cls = django_context.get_model_class_by_fullname(model_fullname)
if model_cls is None:
return ctx.default_return_type
return typecheck_model_method(ctx, django_context, model_cls, 'create')
def get_field_type_from_lookup(ctx: MethodContext, django_context: DjangoContext, model_cls: Type[Model],
*, method: str, lookup: str) -> Optional[MypyType]:
try:
lookup_field = django_context.lookups_context.resolve_lookup(model_cls, lookup)
except FieldError as exc:
ctx.api.fail(exc.args[0], ctx.context)
return None
if isinstance(lookup_field, RelatedField) and lookup_field.column == lookup:
lookup_field = django_context.get_primary_key_field(lookup_field.related_model)
field_get_type = django_context.fields_context.get_field_get_type(helpers.get_typechecker_api(ctx),
lookup_field, method=method)
return field_get_type
def extract_proper_type_queryset_values_list(
ctx: MethodContext, django_context: DjangoContext) -> MypyType:
# called on the Instance, returns QuerySet of something
assert isinstance(ctx.type, Instance)
default_return_type = get_proper_type(ctx.default_return_type)
assert isinstance(default_return_type, Instance)
model_type = _extract_model_type_from_queryset(ctx.type)
if model_type is None:
return AnyType(TypeOfAny.from_omitted_generics)
model_cls = django_context.get_model_class_by_fullname(
model_type.type.fullname)
if model_cls is None:
return default_return_type
flat_expr = helpers.get_call_argument_by_name(ctx, "flat")
if flat_expr is not None and isinstance(flat_expr, NameExpr):
flat = helpers.parse_bool(flat_expr)
else:
flat = False
named_expr = helpers.get_call_argument_by_name(ctx, "named")
if named_expr is not None and isinstance(named_expr, NameExpr):
named = helpers.parse_bool(named_expr)
else:
named = False
if flat and named:
ctx.api.fail("'flat' and 'named' can't be used together", ctx.context)
return helpers.reparametrize_instance(
default_return_type,
[model_type, AnyType(TypeOfAny.from_error)])
# account for possible None
flat = flat or False
named = named or False
is_annotated = is_annotated_model_fullname(model_type.type.fullname)
row_type = get_values_list_row_type(ctx,
django_context,
model_cls,
is_annotated=is_annotated,
flat=flat,
named=named)
return helpers.reparametrize_instance(default_return_type,
[model_type, row_type])
def typecheck_model_method(ctx: Union[FunctionContext, MethodContext],
django_context: DjangoContext,
model_cls: Type[Model], method: str) -> MypyType:
typechecker_api = helpers.get_typechecker_api(ctx)
expected_types = django_context.get_expected_types(typechecker_api,
model_cls,
method=method)
expected_keys = [key for key in expected_types.keys() if key != 'pk']
for actual_name, actual_type in get_actual_types(ctx, expected_keys):
if actual_name not in expected_types:
ctx.api.fail(
'Unexpected attribute "{}" for model "{}"'.format(
actual_name, model_cls.__name__), ctx.context)
continue
helpers.check_types_compatible(
ctx,
expected_type=expected_types[actual_name],
actual_type=actual_type,
error_message='Incompatible type for "{}" of "{}"'.format(
actual_name, model_cls.__name__))
return ctx.default_return_type
def return_proper_field_type_from_get_field(
ctx: MethodContext, django_context: DjangoContext) -> MypyType:
# Options instance
assert isinstance(ctx.type, Instance)
# bail if list of generic params is empty
if len(ctx.type.args) == 0:
return ctx.default_return_type
model_type = ctx.type.args[0]
if not isinstance(model_type, Instance):
return ctx.default_return_type
model_cls = django_context.get_model_class_by_fullname(
model_type.type.fullname)
if model_cls is None:
return ctx.default_return_type
field_name_expr = helpers.get_call_argument_by_name(ctx, 'field_name')
if field_name_expr is None:
return ctx.default_return_type
field_name = helpers.resolve_string_attribute_value(
field_name_expr, django_context)
if field_name is None:
return ctx.default_return_type
try:
field = model_cls._meta.get_field(field_name)
except FieldDoesNotExist as exc:
# if model is abstract, do not raise exception, skip false positives
if not model_cls._meta.abstract:
ctx.api.fail(exc.args[0], ctx.context)
return AnyType(TypeOfAny.from_error)
field_fullname = helpers.get_class_fullname(field.__class__)
return _get_field_instance(ctx, field_fullname)
class NewSemanalDjangoPlugin(Plugin):
def __init__(self, options: Options) -> None:
super().__init__(options)
django_settings_module = extract_django_settings_module(
options.config_file)
self.django_context = DjangoContext(django_settings_module)
def _get_current_queryset_bases(self) -> Dict[str, int]:
model_sym = self.lookup_fully_qualified(
fullnames.QUERYSET_CLASS_FULLNAME)
if model_sym is not None and isinstance(model_sym.node, TypeInfo):
return (helpers.get_django_metadata(model_sym.node).setdefault(
'queryset_bases', {fullnames.QUERYSET_CLASS_FULLNAME: 1}))
else:
return {}
def _get_current_manager_bases(self) -> Dict[str, int]:
model_sym = self.lookup_fully_qualified(
fullnames.MANAGER_CLASS_FULLNAME)
if model_sym is not None and isinstance(model_sym.node, TypeInfo):
return (helpers.get_django_metadata(model_sym.node).setdefault(
'manager_bases', {fullnames.MANAGER_CLASS_FULLNAME: 1}))
else:
return {}
def _get_current_model_bases(self) -> Dict[str, int]:
model_sym = self.lookup_fully_qualified(fullnames.MODEL_CLASS_FULLNAME)
if model_sym is not None and isinstance(model_sym.node, TypeInfo):
return helpers.get_django_metadata(model_sym.node).setdefault(
'model_bases', {fullnames.MODEL_CLASS_FULLNAME: 1})
else:
return {}
def _get_current_form_bases(self) -> Dict[str, int]:
model_sym = self.lookup_fully_qualified(
fullnames.BASEFORM_CLASS_FULLNAME)
if model_sym is not None and isinstance(model_sym.node, TypeInfo):
return (helpers.get_django_metadata(model_sym.node).setdefault(
'baseform_bases', {
fullnames.BASEFORM_CLASS_FULLNAME: 1,
fullnames.FORM_CLASS_FULLNAME: 1,
fullnames.MODELFORM_CLASS_FULLNAME: 1
}))
else:
return {}
def _get_typeinfo_or_none(self, class_name: str) -> Optional[TypeInfo]:
sym = self.lookup_fully_qualified(class_name)
if sym is not None and isinstance(sym.node, TypeInfo):
return sym.node
return None
def _new_dependency(self, module: str) -> Tuple[int, str, int]:
return 10, module, -1
def get_additional_deps(self,
file: MypyFile) -> List[Tuple[int, str, int]]:
# for settings
if file.fullname(
) == 'django.conf' and self.django_context.django_settings_module:
return [
self._new_dependency(
self.django_context.django_settings_module)
]
# for values / values_list
if file.fullname() == 'django.db.models':
return [
self._new_dependency('mypy_extensions'),
self._new_dependency('typing')
]
# for `get_user_model()`
if self.django_context.settings:
if (file.fullname() == 'django.contrib.auth' or file.fullname()
in {'django.http', 'django.http.request'}):
auth_user_model_name = self.django_context.settings.AUTH_USER_MODEL
try:
auth_user_module = self.django_context.apps_registry.get_model(
auth_user_model_name).__module__
except LookupError:
# get_user_model() model app is not installed
return []
return [self._new_dependency(auth_user_module)]
# ensure that all mentioned to='someapp.SomeModel' are loaded with corresponding related Fields
defined_model_classes = self.django_context.model_modules.get(
file.fullname())
if not defined_model_classes:
return []
deps = set()
for model_class in defined_model_classes:
# forward relations
for field in self.django_context.get_model_fields(model_class):
if isinstance(field, RelatedField):
related_model_cls = self.django_context.get_field_related_model_cls(
field)
if related_model_cls is None:
continue
related_model_module = related_model_cls.__module__
if related_model_module != file.fullname():
deps.add(self._new_dependency(related_model_module))
# reverse relations
for relation in model_class._meta.related_objects:
related_model_cls = self.django_context.get_field_related_model_cls(
relation)
related_model_module = related_model_cls.__module__
if related_model_module != file.fullname():
deps.add(self._new_dependency(related_model_module))
return list(deps)
def get_function_hook(
self,
fullname: str) -> Optional[Callable[[FunctionContext], MypyType]]:
if fullname == 'django.contrib.auth.get_user_model':
return partial(settings.get_user_model_hook,
django_context=self.django_context)
manager_bases = self._get_current_manager_bases()
if fullname in manager_bases:
return querysets.determine_proper_manager_type
info = self._get_typeinfo_or_none(fullname)
if info:
if info.has_base(fullnames.FIELD_FULLNAME):
return partial(fields.transform_into_proper_return_type,
django_context=self.django_context)
if helpers.is_model_subclass_info(info, self.django_context):
return partial(init_create.redefine_and_typecheck_model_init,
django_context=self.django_context)
return None
def get_method_hook(
self,
fullname: str) -> Optional[Callable[[MethodContext], MypyType]]:
class_fullname, _, method_name = fullname.rpartition('.')
if method_name == 'get_form_class':
info = self._get_typeinfo_or_none(class_fullname)
if info and info.has_base(fullnames.FORM_MIXIN_CLASS_FULLNAME):
return forms.extract_proper_type_for_get_form_class
if method_name == 'get_form':
info = self._get_typeinfo_or_none(class_fullname)
if info and info.has_base(fullnames.FORM_MIXIN_CLASS_FULLNAME):
return forms.extract_proper_type_for_get_form
if method_name == 'values':
info = self._get_typeinfo_or_none(class_fullname)
if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
return partial(querysets.extract_proper_type_queryset_values,
django_context=self.django_context)
if method_name == 'values_list':
info = self._get_typeinfo_or_none(class_fullname)
if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
return partial(
querysets.extract_proper_type_queryset_values_list,
django_context=self.django_context)
if method_name == 'get_field':
info = self._get_typeinfo_or_none(class_fullname)
if info and info.has_base(fullnames.OPTIONS_CLASS_FULLNAME):
return partial(meta.return_proper_field_type_from_get_field,
django_context=self.django_context)
manager_classes = self._get_current_manager_bases()
if class_fullname in manager_classes and method_name == 'create':
return partial(init_create.redefine_and_typecheck_model_create,
django_context=self.django_context)
if class_fullname in manager_classes and method_name in {
'filter', 'get', 'exclude'
}:
return partial(mypy_django_plugin.transformers.orm_lookups.
typecheck_queryset_filter,
django_context=self.django_context)
return None
def get_base_class_hook(
self,
fullname: str) -> Optional[Callable[[ClassDefContext], None]]:
if (fullname
in self.django_context.all_registered_model_class_fullnames
or fullname in self._get_current_model_bases()):
return partial(transform_model_class,
django_context=self.django_context)
if fullname in self._get_current_manager_bases():
return add_new_manager_base
if fullname in self._get_current_form_bases():
return transform_form_class
return None
def get_attribute_hook(
self,
fullname: str) -> Optional[Callable[[AttributeContext], MypyType]]:
class_name, _, attr_name = fullname.rpartition('.')
if class_name == fullnames.DUMMY_SETTINGS_BASE_CLASS:
return partial(settings.get_type_of_settings_attribute,
django_context=self.django_context)
info = self._get_typeinfo_or_none(class_name)
if info and info.has_base(
fullnames.HTTPREQUEST_CLASS_FULLNAME) and attr_name == 'user':
return partial(
request.set_auth_user_model_as_type_for_request_user,
django_context=self.django_context)
return None
def __init__(self, options: Options) -> None:
super().__init__(options)
django_settings_module = extract_django_settings_module(
options.config_file)
self.django_context = DjangoContext(django_settings_module)
class NewSemanalDjangoPlugin(Plugin):
def __init__(self, options: Options) -> None:
super().__init__(options)
self.plugin_config = DjangoPluginConfig(options.config_file)
# Add paths from MYPYPATH env var
sys.path.extend(mypy_path())
# Add paths from mypy_path config option
sys.path.extend(options.mypy_path)
self.django_context = DjangoContext(
self.plugin_config.django_settings_module)
def _get_current_queryset_bases(self) -> Dict[str, int]:
model_sym = self.lookup_fully_qualified(
fullnames.QUERYSET_CLASS_FULLNAME)
if model_sym is not None and isinstance(model_sym.node, TypeInfo):
return helpers.get_django_metadata(model_sym.node).setdefault(
"queryset_bases", {fullnames.QUERYSET_CLASS_FULLNAME: 1})
else:
return {}
def _get_current_manager_bases(self) -> Dict[str, int]:
model_sym = self.lookup_fully_qualified(
fullnames.MANAGER_CLASS_FULLNAME)
if model_sym is not None and isinstance(model_sym.node, TypeInfo):
return helpers.get_django_metadata(model_sym.node).setdefault(
"manager_bases", {fullnames.MANAGER_CLASS_FULLNAME: 1})
else:
return {}
def _get_current_model_bases(self) -> Dict[str, int]:
model_sym = self.lookup_fully_qualified(fullnames.MODEL_CLASS_FULLNAME)
if model_sym is not None and isinstance(model_sym.node, TypeInfo):
return helpers.get_django_metadata(model_sym.node).setdefault(
"model_bases", {fullnames.MODEL_CLASS_FULLNAME: 1})
else:
return {}
def _get_current_form_bases(self) -> Dict[str, int]:
model_sym = self.lookup_fully_qualified(
fullnames.BASEFORM_CLASS_FULLNAME)
if model_sym is not None and isinstance(model_sym.node, TypeInfo):
return helpers.get_django_metadata(model_sym.node).setdefault(
"baseform_bases",
{
fullnames.BASEFORM_CLASS_FULLNAME: 1,
fullnames.FORM_CLASS_FULLNAME: 1,
fullnames.MODELFORM_CLASS_FULLNAME: 1,
},
)
else:
return {}
def _get_typeinfo_or_none(self, class_name: str) -> Optional[TypeInfo]:
sym = self.lookup_fully_qualified(class_name)
if sym is not None and isinstance(sym.node, TypeInfo):
return sym.node
return None
def _new_dependency(self, module: str) -> Tuple[int, str, int]:
return 10, module, -1
def get_additional_deps(self,
file: MypyFile) -> List[Tuple[int, str, int]]:
# for settings
if file.fullname == "django.conf" and self.django_context.django_settings_module:
return [
self._new_dependency(
self.django_context.django_settings_module)
]
# for values / values_list
if file.fullname == "django.db.models":
return [
self._new_dependency("mypy_extensions"),
self._new_dependency("typing")
]
# for `get_user_model()`
if self.django_context.settings:
if file.fullname == "django.contrib.auth" or file.fullname in {
"django.http", "django.http.request"
}:
auth_user_model_name = self.django_context.settings.AUTH_USER_MODEL
try:
auth_user_module = self.django_context.apps_registry.get_model(
auth_user_model_name).__module__
except LookupError:
# get_user_model() model app is not installed
return []
return [self._new_dependency(auth_user_module)]
# ensure that all mentioned to='someapp.SomeModel' are loaded with corresponding related Fields
defined_model_classes = self.django_context.model_modules.get(
file.fullname)
if not defined_model_classes:
return []
deps = set()
for model_class in defined_model_classes:
# forward relations
for field in self.django_context.get_model_fields(model_class):
if isinstance(field, RelatedField):
related_model_cls = self.django_context.get_field_related_model_cls(
field)
if related_model_cls is None:
continue
related_model_module = related_model_cls.__module__
if related_model_module != file.fullname:
deps.add(self._new_dependency(related_model_module))
# reverse relations
for relation in model_class._meta.related_objects:
related_model_cls = self.django_context.get_field_related_model_cls(
relation)
related_model_module = related_model_cls.__module__
if related_model_module != file.fullname:
deps.add(self._new_dependency(related_model_module))
return list(deps) + [
# for QuerySet.annotate
self._new_dependency("django_stubs_ext"),
# For BaseManager.from_queryset
self._new_dependency("django.db.models.query"),
]
def get_function_hook(
self,
fullname: str) -> Optional[Callable[[FunctionContext], MypyType]]:
if fullname == "django.contrib.auth.get_user_model":
return partial(settings.get_user_model_hook,
django_context=self.django_context)
manager_bases = self._get_current_manager_bases()
if fullname in manager_bases:
return querysets.determine_proper_manager_type
info = self._get_typeinfo_or_none(fullname)
if info:
if info.has_base(fullnames.FIELD_FULLNAME):
return partial(fields.transform_into_proper_return_type,
django_context=self.django_context)
if helpers.is_model_subclass_info(info, self.django_context):
return partial(init_create.redefine_and_typecheck_model_init,
django_context=self.django_context)
return None
def get_method_hook(
self,
fullname: str) -> Optional[Callable[[MethodContext], MypyType]]:
class_fullname, _, method_name = fullname.rpartition(".")
if method_name == "get_form_class":
info = self._get_typeinfo_or_none(class_fullname)
if info and info.has_base(fullnames.FORM_MIXIN_CLASS_FULLNAME):
return forms.extract_proper_type_for_get_form_class
if method_name == "get_form":
info = self._get_typeinfo_or_none(class_fullname)
if info and info.has_base(fullnames.FORM_MIXIN_CLASS_FULLNAME):
return forms.extract_proper_type_for_get_form
manager_classes = self._get_current_manager_bases()
if method_name == "values":
info = self._get_typeinfo_or_none(class_fullname)
if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME
) or class_fullname in manager_classes:
return partial(querysets.extract_proper_type_queryset_values,
django_context=self.django_context)
if method_name == "values_list":
info = self._get_typeinfo_or_none(class_fullname)
if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME
) or class_fullname in manager_classes:
return partial(
querysets.extract_proper_type_queryset_values_list,
django_context=self.django_context)
if method_name == "annotate":
info = self._get_typeinfo_or_none(class_fullname)
if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME
) or class_fullname in manager_classes:
return partial(querysets.extract_proper_type_queryset_annotate,
django_context=self.django_context)
if method_name == "get_field":
info = self._get_typeinfo_or_none(class_fullname)
if info and info.has_base(fullnames.OPTIONS_CLASS_FULLNAME):
return partial(meta.return_proper_field_type_from_get_field,
django_context=self.django_context)
if class_fullname in manager_classes and method_name == "create":
return partial(init_create.redefine_and_typecheck_model_create,
django_context=self.django_context)
if class_fullname in manager_classes and method_name in {
"filter", "get", "exclude"
}:
return partial(
mypy_django_plugin.transformers.orm_lookups.
typecheck_queryset_filter,
django_context=self.django_context,
)
if method_name == "from_queryset":
info = self._get_typeinfo_or_none(class_fullname)
if info and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME):
return fail_if_manager_type_created_in_model_body
return None
def get_base_class_hook(
self,
fullname: str) -> Optional[Callable[[ClassDefContext], None]]:
# Base class is a Model class definition
if (fullname
in self.django_context.all_registered_model_class_fullnames
or fullname in self._get_current_model_bases()):
return partial(transform_model_class,
django_context=self.django_context)
# Base class is a Manager class definition
if fullname in self._get_current_manager_bases():
return add_new_manager_base_hook
# Base class is a Form class definition
if fullname in self._get_current_form_bases():
return transform_form_class
return None
def get_attribute_hook(
self,
fullname: str) -> Optional[Callable[[AttributeContext], MypyType]]:
class_name, _, attr_name = fullname.rpartition(".")
# Lookup of a settings variable
if class_name == fullnames.DUMMY_SETTINGS_BASE_CLASS:
return partial(settings.get_type_of_settings_attribute,
django_context=self.django_context)
info = self._get_typeinfo_or_none(class_name)
# Lookup of the '.is_superuser' attribute
if info and info.has_base(fullnames.PERMISSION_MIXIN_CLASS_FULLNAME
) and attr_name == "is_superuser":
return partial(set_auth_user_model_boolean_fields,
django_context=self.django_context)
# Lookup of the 'request.user' attribute
if info and info.has_base(
fullnames.HTTPREQUEST_CLASS_FULLNAME) and attr_name == "user":
return partial(
request.set_auth_user_model_as_type_for_request_user,
django_context=self.django_context)
# Lookup of the 'user.is_staff' or 'user.is_active' attribute
if info and info.has_base(
fullnames.ABSTRACT_USER_MODEL_FULLNAME) and attr_name in (
"is_staff", "is_active"):
return partial(set_auth_user_model_boolean_fields,
django_context=self.django_context)
# Lookup of a method on a dynamically generated manager class
# i.e. a manager class only existing while mypy is running, not collected from the AST
if (info and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME)
and class_name in self._get_current_manager_bases()):
return resolve_manager_method
return None
def get_type_analyze_hook(
self, fullname: str
) -> Optional[Callable[[AnalyzeTypeContext], MypyType]]:
if fullname in (
"typing.Annotated",
"typing_extensions.Annotated",
"django_stubs_ext.annotations.WithAnnotations",
):
return partial(handle_annotated_type,
django_context=self.django_context)
def get_dynamic_class_hook(
self, fullname: str
) -> Optional[Callable[[DynamicClassDefContext], None]]:
# Create a new manager class definition when a manager's '.from_queryset' classmethod is called
if fullname.endswith("from_queryset"):
class_name, _, _ = fullname.rpartition(".")
info = self._get_typeinfo_or_none(class_name)
if info and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME):
return create_new_manager_class_from_from_queryset_method
return None
def get_values_list_row_type(
ctx: MethodContext,
django_context: DjangoContext,
model_cls: Type[Model],
*,
is_annotated: bool,
flat: bool,
named: bool,
) -> MypyType:
field_lookups = resolve_field_lookups(ctx.args[0], django_context)
if field_lookups is None:
return AnyType(TypeOfAny.from_error)
typechecker_api = helpers.get_typechecker_api(ctx)
if len(field_lookups) == 0:
if flat:
primary_key_field = django_context.get_primary_key_field(model_cls)
lookup_type = get_field_type_from_lookup(
ctx,
django_context,
model_cls,
lookup=primary_key_field.attname,
method="values_list")
assert lookup_type is not None
return lookup_type
elif named:
column_types: "OrderedDict[str, MypyType]" = OrderedDict()
for field in django_context.get_model_fields(model_cls):
column_type = django_context.get_field_get_type(
typechecker_api, field, method="values_list")
column_types[field.attname] = column_type
if is_annotated:
# Return a NamedTuple with a fallback so that it's possible to access any field
return helpers.make_oneoff_named_tuple(
typechecker_api,
"Row",
column_types,
extra_bases=[
typechecker_api.named_generic_type(
ANY_ATTR_ALLOWED_CLASS_FULLNAME, [])
],
)
else:
return helpers.make_oneoff_named_tuple(typechecker_api, "Row",
column_types)
else:
# flat=False, named=False, all fields
if is_annotated:
return typechecker_api.named_generic_type(
"builtins.tuple", [AnyType(TypeOfAny.special_form)])
field_lookups = []
for field in django_context.get_model_fields(model_cls):
field_lookups.append(field.attname)
if len(field_lookups) > 1 and flat:
typechecker_api.fail(
"'flat' is not valid when 'values_list' is called with more than one field",
ctx.context)
return AnyType(TypeOfAny.from_error)
column_types = OrderedDict()
for field_lookup in field_lookups:
lookup_field_type = get_field_type_from_lookup(
ctx,
django_context,
model_cls,
lookup=field_lookup,
method="values_list",
silent_on_error=is_annotated)
if lookup_field_type is None:
if is_annotated:
lookup_field_type = AnyType(TypeOfAny.from_omitted_generics)
else:
return AnyType(TypeOfAny.from_error)
column_types[field_lookup] = lookup_field_type
if flat:
assert len(column_types) == 1
row_type = next(iter(column_types.values()))
elif named:
row_type = helpers.make_oneoff_named_tuple(typechecker_api, "Row",
column_types)
else:
row_type = helpers.make_tuple(typechecker_api,
list(column_types.values()))
return row_type
