VimPlug/django-stubs
1
0
Fork 0
mirror of https://github.com/davidhalter/django-stubs.git synced 2025-12-13 15:31:55 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Adds better id and pk handling (#791)

Browse Source
This commit is contained in:
Nikita Sobolev
2021-12-18 13:09:22 +03:00
committed by GitHub
parent 4a0dd04611
commit b883e192e4
2 changed files with 101 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
mypy_django_plugin/transformers/models.py
View File

@@ -1,7 +1,7 @@
from typing import Dict, List, Optional, Type, Union, cast
from django.db.models.base import Model
from django.db.models.fields import DateField, DateTimeField
from django.db.models.fields import DateField, DateTimeField, Field
from django.db.models.fields.related import ForeignKey
from django.db.models.fields.reverse_related import ManyToManyRel, ManyToOneRel, OneToOneRel
from mypy.checker import TypeChecker
@@ -69,7 +69,7 @@ class ModelClassInitializer:
self.run_with_model_cls(model_cls)
def run_with_model_cls(self, model_cls):
pass
raise NotImplementedError("Implement this in subclasses")
class InjectAnyAsBaseForNestedMeta(ModelClassInitializer):
@@ -95,15 +95,42 @@ class InjectAnyAsBaseForNestedMeta(ModelClassInitializer):
class AddDefaultPrimaryKey(ModelClassInitializer):
def run_with_model_cls(self, model_cls: Type[Model]) -> None:
auto_field = model_cls._meta.auto_field
if auto_field and not self.model_classdef.info.has_readable_member(auto_field.attname):
# autogenerated field
if auto_field:
self.create_autofield(
auto_field=auto_field,
dest_name=auto_field.attname,
existing_field=not self.model_classdef.info.has_readable_member(auto_field.attname),
)
def create_autofield(
self,
auto_field: Field,
dest_name: str,
existing_field: bool,
) -> None:
if existing_field:
auto_field_fullname = helpers.get_class_fullname(auto_field.__class__)
auto_field_info = self.lookup_typeinfo_or_incomplete_defn_error(auto_field_fullname)
set_type, get_type = fields.get_field_descriptor_types(
auto_field_info, is_set_nullable=True, is_get_nullable=False
auto_field_info,
is_set_nullable=True,
is_get_nullable=False,
)
self.add_new_node_to_model_class(dest_name, Instance(auto_field_info, [set_type, get_type]))
class AddPrimaryKeyAlias(AddDefaultPrimaryKey):
def run_with_model_cls(self, model_cls: Type[Model]) -> None:
# We also need to override existing `pk` definition from `stubs`:
auto_field = model_cls._meta.pk
if auto_field:
self.create_autofield(
auto_field=auto_field,
dest_name="pk",
existing_field=self.model_classdef.info.has_readable_member(auto_field.name),
)
self.add_new_node_to_model_class(auto_field.attname, Instance(auto_field_info, [set_type, get_type]))
class AddRelatedModelsId(ModelClassInitializer):
@@ -117,7 +144,7 @@ class AddRelatedModelsId(ModelClassInitializer):
if field_sym is not None and field_sym.node is not None:
error_context = field_sym.node
self.api.fail(
f"Cannot find model {field.related_model!r} " f"referenced in field {field.name!r} ",
f"Cannot find model {field.related_model!r} referenced in field {field.name!r}",
ctx=error_context,
)
self.add_new_node_to_model_class(field.attname, AnyType(TypeOfAny.explicit))
@@ -374,6 +401,7 @@ def process_model_class(ctx: ClassDefContext, django_context: DjangoContext) ->
initializers = [
InjectAnyAsBaseForNestedMeta,
AddDefaultPrimaryKey,
AddPrimaryKeyAlias,
AddRelatedModelsId,
AddManagers,
AddDefaultManagerAttribute,