mirror of
https://github.com/davidhalter/django-stubs.git
synced 2026-05-19 23:09:56 +08:00
new semanal wip 1
This commit is contained in:
Maxim Kurnikov
@@ -18,17 +18,20 @@ from mypy_django_plugin.lib import metadata, fullnames, helpers
|
||||
class ModelClassInitializer(metaclass=ABCMeta):
|
||||
api: SemanticAnalyzerPass2
|
||||
model_classdef: ClassDef
|
||||
app_models_mapping: Optional[Dict[str, str]] = None
|
||||
|
||||
@classmethod
|
||||
def from_ctx(cls, ctx: ClassDefContext):
|
||||
return cls(api=cast(SemanticAnalyzerPass2, ctx.api), model_classdef=ctx.cls)
|
||||
def from_ctx(cls, ctx: ClassDefContext, app_models_mapping: Optional[Dict[str, str]]):
|
||||
return cls(api=cast(SemanticAnalyzerPass2, ctx.api),
|
||||
model_classdef=ctx.cls,
|
||||
app_models_mapping=app_models_mapping)
|
||||
|
||||
def get_meta_attribute(self, name: str) -> Optional[Expression]:
|
||||
meta_node = helpers.get_nested_meta_node_for_current_class(self.model_classdef.info)
|
||||
if meta_node is None:
|
||||
return None
|
||||
|
||||
return helpers.get_assigned_value_for_class(meta_node, name)
|
||||
return helpers.get_assignment_stmt_by_name(meta_node, name)
|
||||
|
||||
def is_abstract_model(self) -> bool:
|
||||
is_abstract_expr = self.get_meta_attribute('abstract')
|
||||
@@ -46,29 +49,74 @@ class ModelClassInitializer(metaclass=ABCMeta):
|
||||
var.is_initialized_in_class = True
|
||||
self.model_classdef.info.names[name] = SymbolTableNode(MDEF, var, plugin_generated=True)
|
||||
|
||||
def model_has_name_defined(self, name: str) -> bool:
|
||||
return name in self.model_classdef.info.names
|
||||
|
||||
@abstractmethod
|
||||
def run(self) -> None:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
def iter_over_one_to_n_related_fields(klass: ClassDef) -> Iterator[Tuple[NameExpr, CallExpr]]:
|
||||
for lvalue, rvalue in helpers.iter_call_assignments(klass):
|
||||
if (isinstance(lvalue, NameExpr)
|
||||
and isinstance(rvalue.callee, MemberExpr)):
|
||||
if rvalue.callee.fullname in {fullnames.FOREIGN_KEY_FULLNAME,
|
||||
fullnames.ONETOONE_FIELD_FULLNAME}:
|
||||
yield lvalue, rvalue
|
||||
for field_name, field_init in helpers.iter_over_field_inits_in_class(klass):
|
||||
field_info = field_init.callee.node
|
||||
assert isinstance(field_info, TypeInfo)
|
||||
|
||||
if helpers.has_any_of_bases(field_init.callee.node, {fullnames.FOREIGN_KEY_FULLNAME,
|
||||
fullnames.ONETOONE_FIELD_FULLNAME}):
|
||||
yield field_name, field_init
|
||||
|
||||
|
||||
class SetIdAttrsForRelatedFields(ModelClassInitializer):
|
||||
class AddReferencesToRelatedModels(ModelClassInitializer):
|
||||
"""
|
||||
For every
|
||||
attr1 = models.ForeignKey(to=MyModel)
|
||||
sets `attr1_id` attribute to the current model.
|
||||
"""
|
||||
|
||||
def run(self) -> None:
|
||||
for lvalue, rvalue in iter_over_one_to_n_related_fields(self.model_classdef):
|
||||
node_name = lvalue.name + '_id'
|
||||
self.add_new_node_to_model_class(name=node_name,
|
||||
typ=self.api.builtin_type('builtins.int'))
|
||||
for field_name, field_init_expr in helpers.iter_over_field_inits_in_class(self.model_classdef):
|
||||
ref_id_name = field_name + '_id'
|
||||
field_info = field_init_expr.callee.node
|
||||
assert isinstance(field_info, TypeInfo)
|
||||
|
||||
if not self.model_has_name_defined(ref_id_name):
|
||||
if helpers.has_any_of_bases(field_info, {fullnames.FOREIGN_KEY_FULLNAME,
|
||||
fullnames.ONETOONE_FIELD_FULLNAME}):
|
||||
self.add_new_node_to_model_class(name=ref_id_name,
|
||||
typ=self.api.builtin_type('builtins.int'))
|
||||
|
||||
# field_init_expr.callee.node
|
||||
#
|
||||
# for field_name, field_init_expr in helpers.iter_call_assignments_in_class(self.model_classdef):
|
||||
# ref_id_name = field_name + '_id'
|
||||
# if not self.model_has_name_defined(ref_id_name):
|
||||
# field_class_info = field_init_expr.callee.node
|
||||
# if not field_class_info:
|
||||
#
|
||||
# if not field_init_expr.callee.node:
|
||||
#
|
||||
# if isinstance(field_init_expr.callee.node, TypeInfo) \
|
||||
# and helpers.has_any_of_bases(field_init_expr.callee.node,
|
||||
# {fullnames.FOREIGN_KEY_FULLNAME,
|
||||
# fullnames.ONETOONE_FIELD_FULLNAME}):
|
||||
# self.add_new_node_to_model_class(name=ref_id_name,
|
||||
# typ=self.api.builtin_type('builtins.int'))
|
||||
|
||||
|
||||
class InjectAnyAsBaseForNestedMeta(ModelClassInitializer):
|
||||
"""
|
||||
Replaces
|
||||
class MyModel(models.Model):
|
||||
class Meta:
|
||||
pass
|
||||
with
|
||||
class MyModel(models.Model):
|
||||
class Meta(Any):
|
||||
pass
|
||||
to get around incompatible Meta inner classes for different models.
|
||||
"""
|
||||
|
||||
def run(self) -> None:
|
||||
meta_node = helpers.get_nested_meta_node_for_current_class(self.model_classdef.info)
|
||||
if meta_node is None:
|
||||
@@ -77,24 +125,24 @@ class InjectAnyAsBaseForNestedMeta(ModelClassInitializer):
|
||||
|
||||
|
||||
class AddDefaultObjectsManager(ModelClassInitializer):
|
||||
def add_new_manager(self, name: str, manager_type: Optional[Instance]) -> None:
|
||||
def _add_new_manager(self, name: str, manager_type: Optional[Instance]) -> None:
|
||||
if manager_type is None:
|
||||
return None
|
||||
self.add_new_node_to_model_class(name, manager_type)
|
||||
|
||||
def add_private_default_manager(self, manager_type: Optional[Instance]) -> None:
|
||||
def _add_private_default_manager(self, manager_type: Optional[Instance]) -> None:
|
||||
if manager_type is None:
|
||||
return None
|
||||
self.add_new_node_to_model_class('_default_manager', manager_type)
|
||||
|
||||
def get_existing_managers(self) -> List[Tuple[str, TypeInfo]]:
|
||||
def _get_existing_managers(self) -> List[Tuple[str, TypeInfo]]:
|
||||
managers = []
|
||||
for base in self.model_classdef.info.mro:
|
||||
for name_expr, member_expr in helpers.iter_call_assignments(base.defn):
|
||||
manager_name = name_expr.name
|
||||
callee_expr = member_expr.callee
|
||||
for manager_name, call_expr in helpers.iter_call_assignments_in_class(base.defn):
|
||||
callee_expr = call_expr.callee
|
||||
if isinstance(callee_expr, IndexExpr):
|
||||
callee_expr = callee_expr.analyzed.expr
|
||||
|
||||
if isinstance(callee_expr, (MemberExpr, NameExpr)) \
|
||||
and isinstance(callee_expr.node, TypeInfo) \
|
||||
and callee_expr.node.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME):
|
||||
@@ -102,12 +150,12 @@ class AddDefaultObjectsManager(ModelClassInitializer):
|
||||
return managers
|
||||
|
||||
def run(self) -> None:
|
||||
existing_managers = self.get_existing_managers()
|
||||
existing_managers = self._get_existing_managers()
|
||||
if existing_managers:
|
||||
first_manager_type = None
|
||||
for manager_name, manager_type_info in existing_managers:
|
||||
manager_type = Instance(manager_type_info, args=[Instance(self.model_classdef.info, [])])
|
||||
self.add_new_manager(name=manager_name, manager_type=manager_type)
|
||||
self._add_new_manager(name=manager_name, manager_type=manager_type)
|
||||
if first_manager_type is None:
|
||||
first_manager_type = manager_type
|
||||
else:
|
||||
@@ -117,33 +165,46 @@ class AddDefaultObjectsManager(ModelClassInitializer):
|
||||
|
||||
first_manager_type = self.api.named_type_or_none(fullnames.MANAGER_CLASS_FULLNAME,
|
||||
args=[Instance(self.model_classdef.info, [])])
|
||||
self.add_new_manager('objects', manager_type=first_manager_type)
|
||||
self._add_new_manager('objects', manager_type=first_manager_type)
|
||||
|
||||
if self.is_abstract_model():
|
||||
return None
|
||||
default_manager_name_expr = self.get_meta_attribute('default_manager_name')
|
||||
if isinstance(default_manager_name_expr, StrExpr):
|
||||
self.add_private_default_manager(self.model_classdef.info.get(default_manager_name_expr.value).type)
|
||||
self._add_private_default_manager(self.model_classdef.info.get(default_manager_name_expr.value).type)
|
||||
else:
|
||||
self.add_private_default_manager(first_manager_type)
|
||||
self._add_private_default_manager(first_manager_type)
|
||||
|
||||
|
||||
class AddIdAttributeIfPrimaryKeyTrueIsNotSet(ModelClassInitializer):
|
||||
class AddDefaultPrimaryKey(ModelClassInitializer):
|
||||
"""
|
||||
Sets default integer `id` attribute, if:
|
||||
* model is not abstract (abstract = False)
|
||||
* there's no field with primary_key=True
|
||||
"""
|
||||
|
||||
def run(self) -> None:
|
||||
if self.is_abstract_model():
|
||||
# no need for .id attr
|
||||
# abstract models cannot be instantiated, and do not need `id` attribute
|
||||
return None
|
||||
|
||||
for _, rvalue in helpers.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')])):
|
||||
for _, field_init_expr in helpers.iter_over_field_inits_in_class(self.model_classdef):
|
||||
if ('primary_key' in field_init_expr.arg_names
|
||||
and self.api.parse_bool(field_init_expr.args[field_init_expr.arg_names.index('primary_key')])):
|
||||
break
|
||||
else:
|
||||
self.add_new_node_to_model_class('id', self.api.builtin_type('builtins.object'))
|
||||
self.add_new_node_to_model_class('id', self.api.builtin_type('builtins.int'))
|
||||
|
||||
|
||||
def _get_to_expr(field_init_expr) -> Expression:
|
||||
if 'to' in field_init_expr.arg_names:
|
||||
return field_init_expr.args[field_init_expr.arg_names.index('to')]
|
||||
else:
|
||||
return field_init_expr.args[0]
|
||||
|
||||
|
||||
class AddRelatedManagers(ModelClassInitializer):
|
||||
def add_related_manager_variable(self, manager_name: str, related_field_type_data: Dict[str, Any]) -> None:
|
||||
def _add_related_manager_variable(self, manager_name: str, related_field_type_data: Dict[str, Any]) -> None:
|
||||
# add dummy related manager for use later
|
||||
self.add_new_node_to_model_class(manager_name, self.api.builtin_type('builtins.object'))
|
||||
|
||||
@@ -153,24 +214,41 @@ class AddRelatedManagers(ModelClassInitializer):
|
||||
|
||||
def run(self) -> None:
|
||||
for module_name, module_file in self.api.modules.items():
|
||||
for model_defn in helpers.iter_over_classdefs(module_file):
|
||||
if not model_defn.info:
|
||||
self.api.defer()
|
||||
for model_classdef in helpers.iter_over_toplevel_classes(module_file):
|
||||
for field_name, field_init in helpers.iter_over_field_inits_in_class(model_classdef):
|
||||
field_info = field_init.callee.node
|
||||
assert isinstance(field_info, TypeInfo)
|
||||
|
||||
for lvalue, field_init in helpers.iter_call_assignments(model_defn):
|
||||
if is_related_field(field_init, module_file):
|
||||
try:
|
||||
referenced_model_fullname = extract_referenced_model_fullname(field_init,
|
||||
module_file=module_file,
|
||||
all_modules=self.api.modules)
|
||||
except helpers.SelfReference:
|
||||
referenced_model_fullname = model_defn.fullname
|
||||
if helpers.has_any_of_bases(field_info, fullnames.RELATED_FIELDS_CLASSES):
|
||||
# try:
|
||||
to_arg_expr = _get_to_expr(field_init)
|
||||
if isinstance(to_arg_expr, NameExpr):
|
||||
referenced_model_fullname = module_file.names[to_arg_expr.name].fullname
|
||||
else:
|
||||
assert isinstance(to_arg_expr, StrExpr)
|
||||
value = to_arg_expr.value
|
||||
if value == 'self':
|
||||
# reference to the same model class
|
||||
referenced_model_fullname = model_classdef.fullname
|
||||
elif '.' not in value:
|
||||
# reference to class in the current module
|
||||
referenced_model_fullname = module_name + '.' + value
|
||||
else:
|
||||
referenced_model_fullname = self.app_models_mapping[value]
|
||||
|
||||
except helpers.SameFileModel as exc:
|
||||
referenced_model_fullname = module_name + '.' + exc.model_cls_name
|
||||
# referenced_model_fullname = extract_referenced_model_fullname(field_init,
|
||||
# module_file=module_file,
|
||||
# all_modules=self.api.modules)
|
||||
# if not referenced_model_fullname:
|
||||
# raise helpers.IncompleteDefnException('Cannot parse referenced model fullname')
|
||||
|
||||
# except helpers.SelfReference:
|
||||
# referenced_model_fullname = model_classdef.fullname
|
||||
#
|
||||
# except helpers.SameFileModel as exc:
|
||||
# referenced_model_fullname = module_name + '.' + exc.model_cls_name
|
||||
|
||||
if self.model_classdef.fullname == referenced_model_fullname:
|
||||
related_name = model_defn.name.lower() + '_set'
|
||||
if 'related_name' in field_init.arg_names:
|
||||
related_name_expr = field_init.args[field_init.arg_names.index('related_name')]
|
||||
if not isinstance(related_name_expr, StrExpr):
|
||||
@@ -180,9 +258,10 @@ class AddRelatedManagers(ModelClassInitializer):
|
||||
if related_name == '+':
|
||||
# No backwards relation is desired
|
||||
continue
|
||||
else:
|
||||
related_name = model_classdef.name.lower() + '_set'
|
||||
|
||||
# Default related_query_name to related_name
|
||||
related_query_name = related_name
|
||||
if 'related_query_name' in field_init.arg_names:
|
||||
related_query_name_expr = field_init.args[field_init.arg_names.index('related_query_name')]
|
||||
if isinstance(related_query_name_expr, StrExpr):
|
||||
@@ -191,20 +270,24 @@ class AddRelatedManagers(ModelClassInitializer):
|
||||
# not string 'related_query_name=' is not yet supported
|
||||
related_query_name = None
|
||||
# TODO: Handle defaulting to model name if related_name is not set
|
||||
|
||||
# as long as Model is not a Generic, one level depth is fine
|
||||
if field_init.callee.name in {'ForeignKey', 'ManyToManyField'}:
|
||||
field_type_data = {
|
||||
'manager': fullnames.RELATED_MANAGER_CLASS_FULLNAME,
|
||||
'of': [model_defn.info.fullname()]
|
||||
}
|
||||
else:
|
||||
field_type_data = {
|
||||
'manager': model_defn.info.fullname(),
|
||||
'of': []
|
||||
}
|
||||
related_query_name = related_name
|
||||
|
||||
self.add_related_manager_variable(related_name, related_field_type_data=field_type_data)
|
||||
# if helpers.has_any_of_bases(field_info, {fullnames.FOREIGN_KEY_FULLNAME,
|
||||
# fullnames.MANYTOMANY_FIELD_FULLNAME}):
|
||||
# # as long as Model is not a Generic, one level depth is fine
|
||||
# field_type_data = {
|
||||
# 'manager': fullnames.RELATED_MANAGER_CLASS_FULLNAME,
|
||||
# 'of': [model_classdef.info.fullname()]
|
||||
# }
|
||||
# else:
|
||||
# field_type_data = {
|
||||
# 'manager': model_classdef.info.fullname(),
|
||||
# 'of': []
|
||||
# }
|
||||
self.add_new_node_to_model_class(related_name, self.api.builtin_type('builtins.object'))
|
||||
|
||||
# self._add_related_manager_variable(related_name, related_field_type_data=field_type_data)
|
||||
|
||||
if related_query_name is not None:
|
||||
# Only create related_query_name if it is a string literal
|
||||
@@ -239,22 +322,24 @@ def is_related_field(expr: CallExpr, module_file: MypyFile) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def extract_referenced_model_fullname(rvalue_expr: CallExpr,
|
||||
def extract_referenced_model_fullname(field_init_expr: CallExpr,
|
||||
module_file: MypyFile,
|
||||
all_modules: Dict[str, MypyFile]) -> Optional[str]:
|
||||
""" Returns fullname of a Model referenced in "to=" argument of the CallExpr"""
|
||||
if 'to' in rvalue_expr.arg_names:
|
||||
to_expr = rvalue_expr.args[rvalue_expr.arg_names.index('to')]
|
||||
if 'to' in field_init_expr.arg_names:
|
||||
to_expr = field_init_expr.args[field_init_expr.arg_names.index('to')]
|
||||
else:
|
||||
to_expr = rvalue_expr.args[0]
|
||||
to_expr = field_init_expr.args[0]
|
||||
|
||||
if isinstance(to_expr, NameExpr):
|
||||
return module_file.names[to_expr.name].fullname
|
||||
|
||||
elif isinstance(to_expr, StrExpr):
|
||||
typ_fullname = helpers.get_model_fullname_from_string(to_expr.value, all_modules)
|
||||
if typ_fullname is None:
|
||||
return None
|
||||
return typ_fullname
|
||||
|
||||
return None
|
||||
|
||||
|
||||
@@ -284,16 +369,18 @@ def add_get_set_attr_fallback_to_any(ctx: ClassDefContext):
|
||||
add_method(ctx, '__setattr__', [name_arg, value_arg], any)
|
||||
|
||||
|
||||
def process_model_class(ctx: ClassDefContext, ignore_unknown_attributes: bool) -> None:
|
||||
def process_model_class(ctx: ClassDefContext,
|
||||
ignore_unknown_attributes: bool,
|
||||
app_models_mapping: Optional[Dict[str, str]]) -> None:
|
||||
initializers = [
|
||||
InjectAnyAsBaseForNestedMeta,
|
||||
AddDefaultPrimaryKey,
|
||||
AddReferencesToRelatedModels,
|
||||
AddDefaultObjectsManager,
|
||||
AddIdAttributeIfPrimaryKeyTrueIsNotSet,
|
||||
SetIdAttrsForRelatedFields,
|
||||
AddRelatedManagers,
|
||||
]
|
||||
for initializer_cls in initializers:
|
||||
initializer_cls.from_ctx(ctx).run()
|
||||
initializer_cls.from_ctx(ctx, app_models_mapping).run()
|
||||
|
||||
add_dummy_init_method(ctx)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user