add some support for proxy models

This commit is contained in:
Maxim Kurnikov
2019-07-22 21:48:59 +03:00
parent 57796077c6
commit df021f6baa
6 changed files with 61 additions and 29 deletions

View File

@@ -1,12 +1,16 @@
from typing import Any, Optional, Tuple, List, overload from typing import Any, Optional, Tuple, List, overload, TypeVar
from django.db.models.base import Model
from django.db import models from django.db import models
class BaseUserManager(models.Manager): _T = TypeVar('_T', bound=Model)
class BaseUserManager(models.Manager[_T]):
@classmethod @classmethod
def normalize_email(cls, email: Optional[str]) -> str: ... def normalize_email(cls, email: Optional[str]) -> str: ...
def make_random_password(self, length: int = ..., allowed_chars: str = ...) -> str: ... def make_random_password(self, length: int = ..., allowed_chars: str = ...) -> str: ...
def get_by_natural_key(self, username: Optional[str]) -> AbstractBaseUser: ... def get_by_natural_key(self, username: Optional[str]) -> _T: ...
class AbstractBaseUser(models.Model): class AbstractBaseUser(models.Model):
password: models.CharField = ... password: models.CharField = ...

View File

@@ -1,7 +1,8 @@
from typing import Any, Collection, Optional, Set, Tuple, Type, Union from typing import Any, List, Optional, Set, Tuple, Type, Union, TypeVar
from django.contrib.auth.base_user import AbstractBaseUser as AbstractBaseUser, BaseUserManager as BaseUserManager from django.contrib.auth.base_user import AbstractBaseUser as AbstractBaseUser, BaseUserManager as BaseUserManager
from django.contrib.contenttypes.models import ContentType from django.contrib.contenttypes.models import ContentType
from django.db.models.base import Model
from django.db.models.manager import EmptyManager from django.db.models.manager import EmptyManager
from django.contrib.auth.validators import UnicodeUsernameValidator from django.contrib.auth.validators import UnicodeUsernameValidator
@@ -27,13 +28,15 @@ class Group(models.Model):
permissions: models.ManyToManyField = models.ManyToManyField(Permission) permissions: models.ManyToManyField = models.ManyToManyField(Permission)
def natural_key(self): ... def natural_key(self): ...
class UserManager(BaseUserManager): _T = TypeVar('_T', bound=Model)
class UserManager(BaseUserManager[_T]):
def create_user( def create_user(
self, username: str, email: Optional[str] = ..., password: Optional[str] = ..., **extra_fields: Any self, username: str, email: Optional[str] = ..., password: Optional[str] = ..., **extra_fields: Any
) -> AbstractBaseUser: ... ) -> _T: ...
def create_superuser( def create_superuser(
self, username: str, email: Optional[str], password: Optional[str], **extra_fields: Any self, username: str, email: Optional[str], password: Optional[str], **extra_fields: Any
) -> AbstractBaseUser: ... ) -> _T: ...
class PermissionsMixin(models.Model): class PermissionsMixin(models.Model):
is_superuser: models.BooleanField = ... is_superuser: models.BooleanField = ...

View File

@@ -1,22 +1,19 @@
import os import os
from collections import defaultdict from collections import defaultdict
from contextlib import contextmanager from contextlib import contextmanager
from typing import ( from typing import (Any, Dict, Iterator, List, Optional, TYPE_CHECKING, Tuple, Type)
TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple, Type,
)
from django.contrib.postgres.fields import ArrayField
from django.core.exceptions import FieldError from django.core.exceptions import FieldError
from django.db.models.base import Model from django.db.models.base import Model
from django.db.models.fields import CharField, Field, AutoField
from django.db.models.fields.related import ForeignKey, RelatedField from django.db.models.fields.related import ForeignKey, RelatedField
from django.db.models.fields.reverse_related import ForeignObjectRel from django.db.models.fields.reverse_related import ForeignObjectRel
from django.db.models.sql.query import Query from django.db.models.sql.query import Query
from django.utils.functional import cached_property from django.utils.functional import cached_property
from mypy.checker import TypeChecker from mypy.checker import TypeChecker
from mypy.types import Instance from mypy.types import Instance, Type as MypyType
from mypy.types import Type as MypyType
from django.contrib.postgres.fields import ArrayField
from django.db.models.fields import AutoField, CharField, Field
from mypy_django_plugin.lib import helpers from mypy_django_plugin.lib import helpers
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -210,13 +207,19 @@ class DjangoContext:
if isinstance(field, ForeignKey): if isinstance(field, ForeignKey):
field_name = field.name field_name = field.name
foreign_key_info = helpers.lookup_class_typeinfo(api, field.__class__) foreign_key_info = helpers.lookup_class_typeinfo(api, field.__class__)
related_model_info = helpers.lookup_class_typeinfo(api, field.related_model)
related_model = field.related_model
if related_model._meta.proxy_for_model:
related_model = field.related_model._meta.proxy_for_model
related_model_info = helpers.lookup_class_typeinfo(api, related_model)
is_nullable = self.fields_context.get_field_nullability(field, method) is_nullable = self.fields_context.get_field_nullability(field, method)
foreign_key_set_type = helpers.get_private_descriptor_type(foreign_key_info, foreign_key_set_type = helpers.get_private_descriptor_type(foreign_key_info,
'_pyi_private_set_type', '_pyi_private_set_type',
is_nullable=is_nullable) is_nullable=is_nullable)
model_set_type = helpers.convert_any_to_type(foreign_key_set_type, model_set_type = helpers.convert_any_to_type(foreign_key_set_type,
Instance(related_model_info, [])) Instance(related_model_info, []))
expected_types[field_name] = model_set_type expected_types[field_name] = model_set_type
elif isinstance(field, GenericForeignKey): elif isinstance(field, GenericForeignKey):

View File

@@ -1,17 +1,13 @@
from collections import OrderedDict from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union from typing import Any, Dict, List, Optional, Set, TYPE_CHECKING, Union
from mypy import checker from mypy import checker
from mypy.checker import TypeChecker from mypy.checker import TypeChecker
from mypy.mro import calculate_mro from mypy.mro import calculate_mro
from mypy.nodes import ( from mypy.nodes import (Block, ClassDef, Expression, GDEF, MDEF, MemberExpr, MypyFile, NameExpr, StrExpr, SymbolNode,
GDEF, MDEF, Block, ClassDef, Expression, MemberExpr, MypyFile, NameExpr, StrExpr, SymbolNode, SymbolTable, SymbolTable, SymbolTableNode, TypeInfo, Var)
SymbolTableNode, TypeInfo, Var,
)
from mypy.plugin import CheckerPluginInterface, FunctionContext, MethodContext from mypy.plugin import CheckerPluginInterface, FunctionContext, MethodContext
from mypy.types import AnyType, Instance, NoneTyp, TupleType from mypy.types import AnyType, Instance, NoneTyp, TupleType, Type as MypyType, TypeOfAny, TypedDictType, UnionType
from mypy.types import Type as MypyType
from mypy.types import TypedDictType, TypeOfAny, UnionType
if TYPE_CHECKING: if TYPE_CHECKING:
from mypy_django_plugin.django.context import DjangoContext from mypy_django_plugin.django.context import DjangoContext

View File

@@ -39,15 +39,20 @@ def fill_descriptor_types_for_related_field(ctx: FunctionContext, django_context
return AnyType(TypeOfAny.from_error) return AnyType(TypeOfAny.from_error)
assert isinstance(current_field, RelatedField) assert isinstance(current_field, RelatedField)
referred_to_typeinfo = helpers.lookup_class_typeinfo(ctx.api, current_field.related_model)
referred_to_type = Instance(referred_to_typeinfo, []) related_model = related_model_to_set = current_field.related_model
if related_model_to_set._meta.proxy_for_model:
related_model_to_set = related_model._meta.proxy_for_model
related_model_info = helpers.lookup_class_typeinfo(ctx.api, related_model)
related_model_to_set_info = helpers.lookup_class_typeinfo(ctx.api, related_model_to_set)
default_related_field_type = set_descriptor_types_for_field(ctx) default_related_field_type = set_descriptor_types_for_field(ctx)
# replace Any with referred_to_type # replace Any with referred_to_type
args = [] args = [
for default_arg in default_related_field_type.args: helpers.convert_any_to_type(default_related_field_type.args[0], Instance(related_model_to_set_info, [])),
args.append(helpers.convert_any_to_type(default_arg, referred_to_type)) helpers.convert_any_to_type(default_related_field_type.args[1], Instance(related_model_info, [])),
]
return helpers.reparametrize_instance(default_related_field_type, new_args=args) return helpers.reparametrize_instance(default_related_field_type, new_args=args)

View File

@@ -0,0 +1,21 @@
- case: foreign_key_to_proxy_model_accepts_first_non_proxy_model
main: |
from myapp.models import Blog, Publisher, PublisherProxy
Blog(publisher=Publisher())
Blog.objects.create(publisher=Publisher())
Blog().publisher = Publisher()
reveal_type(Blog().publisher) # N: Revealed type is 'myapp.models.PublisherProxy*'
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
class Publisher(models.Model):
pass
class PublisherProxy(Publisher):
class Meta:
proxy = True
class Blog(models.Model):
publisher = models.ForeignKey(to=PublisherProxy, on_delete=models.CASCADE)