mirror of
https://github.com/davidhalter/django-stubs.git
synced 2025-12-09 05:24:53 +08:00
add support for managers as generics
This commit is contained in:
Maxim Kurnikov
@@ -1,3 +1,5 @@
|
||||
import dataclasses
|
||||
from abc import abstractmethod, ABCMeta
|
||||
from typing import cast, Iterator, Tuple, Optional, Dict
|
||||
|
||||
from mypy.nodes import ClassDef, AssignmentStmt, CallExpr, MemberExpr, StrExpr, NameExpr, MDEF, TypeInfo, Var, SymbolTableNode, \
|
||||
@@ -9,6 +11,47 @@ from mypy.types import Instance
|
||||
from mypy_django_plugin import helpers
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ModelClassInitializer(metaclass=ABCMeta):
|
||||
api: SemanticAnalyzerPass2
|
||||
model_classdef: ClassDef
|
||||
|
||||
@classmethod
|
||||
def from_ctx(cls, ctx: ClassDefContext):
|
||||
return cls(api=cast(SemanticAnalyzerPass2, ctx.api), model_classdef=ctx.cls)
|
||||
|
||||
def get_nested_meta_node(self) -> Optional[TypeInfo]:
|
||||
metaclass_sym = self.model_classdef.info.names.get('Meta')
|
||||
if metaclass_sym is not None and isinstance(metaclass_sym.node, TypeInfo):
|
||||
return metaclass_sym.node
|
||||
return None
|
||||
|
||||
def is_abstract_model(self) -> bool:
|
||||
meta_node = self.get_nested_meta_node()
|
||||
if meta_node is None:
|
||||
return False
|
||||
|
||||
for lvalue, rvalue in iter_over_assignments(meta_node.defn):
|
||||
if isinstance(lvalue, NameExpr) and lvalue.name == 'abstract':
|
||||
is_abstract = self.api.parse_bool(rvalue)
|
||||
if is_abstract:
|
||||
# abstract model do not need 'objects' queryset
|
||||
return True
|
||||
return False
|
||||
|
||||
def add_new_node_to_model_class(self, name: str, typ: Instance) -> None:
|
||||
var = Var(name=name, type=typ)
|
||||
var.info = typ.type
|
||||
var._fullname = self.model_classdef.info.fullname() + '.' + name
|
||||
var.is_inferred = True
|
||||
var.is_initialized_in_class = True
|
||||
self.model_classdef.info.names[name] = SymbolTableNode(MDEF, var)
|
||||
|
||||
@abstractmethod
|
||||
def run(self) -> None:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
def add_new_var_node_to_class(class_type: TypeInfo, name: str, typ: Instance) -> None:
|
||||
var = Var(name=name, type=typ)
|
||||
var.info = typ.type
|
||||
@@ -37,78 +80,76 @@ def iter_call_assignments(klass: ClassDef) -> Iterator[Tuple[Lvalue, CallExpr]]:
|
||||
def iter_over_one_to_n_related_fields(klass: ClassDef) -> Iterator[Tuple[NameExpr, CallExpr]]:
|
||||
for lvalue, rvalue in iter_call_assignments(klass):
|
||||
if (isinstance(lvalue, NameExpr)
|
||||
and isinstance(rvalue.callee, MemberExpr)):
|
||||
and isinstance(rvalue.callee, MemberExpr)):
|
||||
if rvalue.callee.fullname in {helpers.FOREIGN_KEY_FULLNAME,
|
||||
helpers.ONETOONE_FIELD_FULLNAME}:
|
||||
yield lvalue, rvalue
|
||||
|
||||
|
||||
def get_nested_meta_class(model_type: TypeInfo) -> Optional[TypeInfo]:
|
||||
metaclass_sym = model_type.names.get('Meta')
|
||||
if metaclass_sym is not None and isinstance(metaclass_sym.node, TypeInfo):
|
||||
return metaclass_sym.node
|
||||
return None
|
||||
class SetIdAttrsForRelatedFields(ModelClassInitializer):
|
||||
def run(self) -> None:
|
||||
for lvalue, rvalue in iter_over_one_to_n_related_fields(self.model_classdef):
|
||||
self.add_new_node_to_model_class(lvalue.name + '_id',
|
||||
typ=self.api.named_type('__builtins__.int'))
|
||||
|
||||
|
||||
def is_abstract_model(ctx: ClassDefContext) -> bool:
|
||||
meta_node = get_nested_meta_class(ctx.cls.info)
|
||||
if meta_node is None:
|
||||
return False
|
||||
|
||||
for lvalue, rvalue in iter_over_assignments(meta_node.defn):
|
||||
if isinstance(lvalue, NameExpr) and lvalue.name == 'abstract':
|
||||
is_abstract = ctx.api.parse_bool(rvalue)
|
||||
if is_abstract:
|
||||
# abstract model do not need 'objects' queryset
|
||||
return True
|
||||
return False
|
||||
class InjectAnyAsBaseForNestedMeta(ModelClassInitializer):
|
||||
def run(self) -> None:
|
||||
meta_node = self.get_nested_meta_node()
|
||||
if meta_node is None:
|
||||
return None
|
||||
meta_node.fallback_to_any = True
|
||||
|
||||
|
||||
def set_fieldname_attrs_for_related_fields(ctx: ClassDefContext) -> None:
|
||||
api = ctx.api
|
||||
for lvalue, rvalue in iter_over_one_to_n_related_fields(ctx.cls):
|
||||
property_name = lvalue.name + '_id'
|
||||
add_new_var_node_to_class(ctx.cls.info, property_name,
|
||||
typ=api.named_type('__builtins__.int'))
|
||||
class AddDefaultObjectsManager(ModelClassInitializer):
|
||||
def run(self) -> None:
|
||||
if 'objects' in self.model_classdef.info.names:
|
||||
return None
|
||||
|
||||
if self.is_abstract_model():
|
||||
# abstract models do not need 'objects' queryset
|
||||
return None
|
||||
|
||||
typ = self.api.named_type_or_none(helpers.MANAGER_CLASS_FULLNAME,
|
||||
args=[Instance(self.model_classdef.info, [])])
|
||||
if not typ:
|
||||
return None
|
||||
self.add_new_node_to_model_class('objects', typ)
|
||||
|
||||
|
||||
def add_int_id_attribute_if_primary_key_true_is_not_present(ctx: ClassDefContext) -> None:
|
||||
api = cast(SemanticAnalyzerPass2, ctx.api)
|
||||
if is_abstract_model(ctx):
|
||||
return None
|
||||
class AddIdAttributeIfPrimaryKeyTrueIsNotSet(ModelClassInitializer):
|
||||
def run(self) -> None:
|
||||
if self.is_abstract_model():
|
||||
# no need for .id attr
|
||||
return None
|
||||
|
||||
for _, rvalue in iter_call_assignments(ctx.cls):
|
||||
if ('primary_key' in rvalue.arg_names and
|
||||
api.parse_bool(rvalue.args[rvalue.arg_names.index('primary_key')])):
|
||||
break
|
||||
else:
|
||||
add_new_var_node_to_class(ctx.cls.info, 'id', api.builtin_type('builtins.int'))
|
||||
for _, rvalue in iter_call_assignments(self.model_classdef):
|
||||
if ('primary_key' in rvalue.arg_names
|
||||
and self.api.parse_bool(rvalue.args[rvalue.arg_names.index('primary_key')])):
|
||||
break
|
||||
else:
|
||||
self.add_new_node_to_model_class('id', self.api.builtin_type('builtins.int'))
|
||||
|
||||
|
||||
def set_objects_queryset_to_model_class(ctx: ClassDefContext) -> None:
|
||||
# search over mro
|
||||
objects_sym = ctx.cls.info.get('objects')
|
||||
if objects_sym is not None:
|
||||
return None
|
||||
|
||||
# only direct Meta class
|
||||
if is_abstract_model(ctx):
|
||||
# abstract model do not need 'objects' queryset
|
||||
return None
|
||||
|
||||
api = cast(SemanticAnalyzerPass2, ctx.api)
|
||||
typ = api.named_type_or_none(helpers.QUERYSET_CLASS_FULLNAME,
|
||||
args=[Instance(ctx.cls.info, [])])
|
||||
if not typ:
|
||||
return None
|
||||
add_new_var_node_to_class(ctx.cls.info, 'objects', typ=typ)
|
||||
|
||||
|
||||
def inject_any_as_base_for_nested_class_meta(ctx: ClassDefContext) -> None:
|
||||
meta_node = get_nested_meta_class(ctx.cls.info)
|
||||
if meta_node is None:
|
||||
return None
|
||||
meta_node.fallback_to_any = True
|
||||
class AddRelatedManagers(ModelClassInitializer):
|
||||
def run(self) -> None:
|
||||
for module_name, module_file in self.api.modules.items():
|
||||
for defn in iter_over_classdefs(module_file):
|
||||
for lvalue, rvalue in iter_call_assignments(defn):
|
||||
if is_related_field(rvalue, module_file):
|
||||
ref_to_fullname = extract_ref_to_fullname(rvalue,
|
||||
module_file=module_file,
|
||||
all_modules=self.api.modules)
|
||||
if self.model_classdef.fullname == ref_to_fullname:
|
||||
if 'related_name' in rvalue.arg_names:
|
||||
related_name_expr = rvalue.args[rvalue.arg_names.index('related_name')]
|
||||
if not isinstance(related_name_expr, StrExpr):
|
||||
return None
|
||||
related_name = related_name_expr.value
|
||||
typ = get_related_field_type(rvalue, self.api, defn.info)
|
||||
if typ is None:
|
||||
return None
|
||||
self.add_new_node_to_model_class(related_name, typ)
|
||||
|
||||
|
||||
def iter_over_classdefs(module_file: MypyFile) -> Iterator[ClassDef]:
|
||||
@@ -119,8 +160,8 @@ def iter_over_classdefs(module_file: MypyFile) -> Iterator[ClassDef]:
|
||||
|
||||
def get_related_field_type(rvalue: CallExpr, api: SemanticAnalyzerPass2,
|
||||
related_model_typ: TypeInfo) -> Optional[Instance]:
|
||||
if rvalue.callee.name == 'ForeignKey':
|
||||
return api.named_type_or_none(helpers.QUERYSET_CLASS_FULLNAME,
|
||||
if rvalue.callee.name in {'ForeignKey', 'ManyToManyField'}:
|
||||
return api.named_type_or_none(helpers.RELATED_MANAGER_CLASS_FULLNAME,
|
||||
args=[Instance(related_model_typ, [])])
|
||||
else:
|
||||
return Instance(related_model_typ, [])
|
||||
@@ -129,7 +170,9 @@ def get_related_field_type(rvalue: CallExpr, api: SemanticAnalyzerPass2,
|
||||
def is_related_field(expr: CallExpr, module_file: MypyFile) -> bool:
|
||||
if isinstance(expr.callee, MemberExpr) and isinstance(expr.callee.expr, NameExpr):
|
||||
module = module_file.names[expr.callee.expr.name]
|
||||
if module.fullname == 'django.db.models' and expr.callee.name in {'ForeignKey', 'OneToOneField'}:
|
||||
if module.fullname == 'django.db.models' and expr.callee.name in {'ForeignKey',
|
||||
'OneToOneField',
|
||||
'ManyToManyField'}:
|
||||
return True
|
||||
return False
|
||||
|
||||
@@ -150,31 +193,16 @@ def extract_ref_to_fullname(rvalue_expr: CallExpr,
|
||||
return None
|
||||
|
||||
|
||||
def add_related_managers(ctx: ClassDefContext):
|
||||
api = cast(SemanticAnalyzerPass2, ctx.api)
|
||||
for module_name, module_file in ctx.api.modules.items():
|
||||
for defn in iter_over_classdefs(module_file):
|
||||
for lvalue, rvalue in iter_call_assignments(defn):
|
||||
if is_related_field(rvalue, module_file):
|
||||
ref_to_fullname = extract_ref_to_fullname(rvalue, module_file=module_file,
|
||||
all_modules=api.modules)
|
||||
if ctx.cls.fullname == ref_to_fullname:
|
||||
if 'related_name' in rvalue.arg_names:
|
||||
related_name_expr = rvalue.args[rvalue.arg_names.index('related_name')]
|
||||
if not isinstance(related_name_expr, StrExpr):
|
||||
return None
|
||||
related_name = related_name_expr.value
|
||||
typ = get_related_field_type(rvalue, api, defn.info)
|
||||
if typ is None:
|
||||
return None
|
||||
add_new_var_node_to_class(ctx.cls.info, related_name, typ)
|
||||
|
||||
|
||||
def process_model_class(ctx: ClassDefContext) -> None:
|
||||
add_related_managers(ctx)
|
||||
inject_any_as_base_for_nested_class_meta(ctx)
|
||||
set_fieldname_attrs_for_related_fields(ctx)
|
||||
add_int_id_attribute_if_primary_key_true_is_not_present(ctx)
|
||||
set_objects_queryset_to_model_class(ctx)
|
||||
initializers = [
|
||||
InjectAnyAsBaseForNestedMeta,
|
||||
AddDefaultObjectsManager,
|
||||
AddIdAttributeIfPrimaryKeyTrueIsNotSet,
|
||||
SetIdAttrsForRelatedFields,
|
||||
AddRelatedManagers
|
||||
]
|
||||
for initializer_cls in initializers:
|
||||
initializer_cls.from_ctx(ctx).run()
|
||||
|
||||
# allow unspecified attributes for now
|
||||
ctx.cls.info.fallback_to_any = True
|
||||
|
||||
Reference in New Issue
Block a user