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,22 +1,19 @@
import os
from collections import defaultdict
from contextlib import contextmanager
from typing import (
TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple, Type,
)
from typing import (Any, Dict, Iterator, List, Optional, TYPE_CHECKING, Tuple, Type)
from django.contrib.postgres.fields import ArrayField
from django.core.exceptions import FieldError
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.reverse_related import ForeignObjectRel
from django.db.models.sql.query import Query
from django.utils.functional import cached_property
from mypy.checker import TypeChecker
from mypy.types import Instance
from mypy.types import Type as MypyType
from mypy.types import Instance, 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
if TYPE_CHECKING:
@@ -210,13 +207,19 @@ class DjangoContext:
if isinstance(field, ForeignKey):
field_name = field.name
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)
foreign_key_set_type = helpers.get_private_descriptor_type(foreign_key_info,
'_pyi_private_set_type',
is_nullable=is_nullable)
model_set_type = helpers.convert_any_to_type(foreign_key_set_type,
Instance(related_model_info, []))
expected_types[field_name] = model_set_type
elif isinstance(field, GenericForeignKey):

View File

@@ -1,17 +1,13 @@
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.checker import TypeChecker
from mypy.mro import calculate_mro
from mypy.nodes import (
GDEF, MDEF, Block, ClassDef, Expression, MemberExpr, MypyFile, NameExpr, StrExpr, SymbolNode, SymbolTable,
SymbolTableNode, TypeInfo, Var,
)
from mypy.nodes import (Block, ClassDef, Expression, GDEF, MDEF, MemberExpr, MypyFile, NameExpr, StrExpr, SymbolNode,
SymbolTable, SymbolTableNode, TypeInfo, Var)
from mypy.plugin import CheckerPluginInterface, FunctionContext, MethodContext
from mypy.types import AnyType, Instance, NoneTyp, TupleType
from mypy.types import Type as MypyType
from mypy.types import TypedDictType, TypeOfAny, UnionType
from mypy.types import AnyType, Instance, NoneTyp, TupleType, Type as MypyType, TypeOfAny, TypedDictType, UnionType
if TYPE_CHECKING:
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)
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)
# replace Any with referred_to_type
args = []
for default_arg in default_related_field_type.args:
args.append(helpers.convert_any_to_type(default_arg, referred_to_type))
args = [
helpers.convert_any_to_type(default_related_field_type.args[0], Instance(related_model_to_set_info, [])),
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)