This commit is contained in:
Maxim Kurnikov
2020-03-15 00:58:11 +03:00
parent 0b1507c81e
commit 1419b144d9
20 changed files with 513 additions and 321 deletions

View File

@@ -4,11 +4,17 @@ from mypy.checker import gen_unique_name
from mypy.nodes import NameExpr, TypeInfo, SymbolTableNode, StrExpr
from mypy.types import Type as MypyType, TypeVarType, TypeVarDef, Instance
from mypy_django_plugin.lib import helpers, fullnames, chk_helpers, sem_helpers
from mypy_django_plugin.lib import helpers, fullnames
from mypy_django_plugin.transformers.managers import iter_all_custom_queryset_methods
class CreateNewManagerClassFrom_FromQuerySet(helpers.DynamicClassPluginCallback):
def set_manager_mapping(self, runtime_manager_fullname: str, generated_manager_fullname: str) -> None:
base_model_info = self.lookup_typeinfo_or_defer(fullnames.MODEL_CLASS_FULLNAME)
assert base_model_info is not None
managers_metadata = base_model_info.metadata.setdefault('managers', {})
managers_metadata[runtime_manager_fullname] = generated_manager_fullname
def create_typevar_in_current_module(self, name: str,
upper_bound: Optional[MypyType] = None) -> TypeVarDef:
tvar_name = gen_unique_name(name, self.semanal_api.globals)
@@ -48,19 +54,20 @@ class CreateNewManagerClassFrom_FromQuerySet(helpers.DynamicClassPluginCallback)
parent_manager_type = Instance(callee_manager_info, [model_tvar_type])
# instantiate with a proper model, Manager[MyModel], filling all Manager type vars in process
queryset_type = Instance(passed_queryset_info, [Instance(base_model_info, [])])
new_manager_info = self.new_typeinfo(self.class_name,
bases=[parent_manager_type])
bases=[queryset_type, parent_manager_type])
new_manager_info.defn.type_vars = [model_tvar_defn]
new_manager_info.type_vars = [model_tvar_defn.name]
new_manager_info.set_line(self.call_expr)
# copy methods from passed_queryset_info with self type replaced
self_type = Instance(new_manager_info, [model_tvar_type])
for name, method_node in iter_all_custom_queryset_methods(passed_queryset_info):
self.add_method_from_signature(method_node,
name,
self_type,
new_manager_info.defn)
# self_type = Instance(new_manager_info, [model_tvar_type])
# for name, method_node in iter_all_custom_queryset_methods(passed_queryset_info):
# self.add_method_from_signature(method_node,
# name,
# self_type,
# new_manager_info.defn)
new_manager_sym = SymbolTableNode(self.semanal_api.current_symbol_kind(),
new_manager_info,
@@ -75,5 +82,5 @@ class CreateNewManagerClassFrom_FromQuerySet(helpers.DynamicClassPluginCallback)
runtime_manager_class_name = class_name_arg.value
new_manager_name = runtime_manager_class_name or (callee_manager_info.name + 'From' + queryset_class_name)
django_generated_manager_name = 'django.db.models.manager.' + new_manager_name
base_model_info.metadata.setdefault('managers', {})[django_generated_manager_name] = new_manager_info.fullname
self.set_manager_mapping(f'django.db.models.manager.{new_manager_name}',
new_manager_info.fullname)