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