Esempio n. 1
0
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
Esempio n. 2
0
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)
Esempio n. 3
0
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
Esempio n. 4
0
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