optimize hooks a bit

This commit is contained in:
Maxim Kurnikov
2019-02-22 00:12:23 +03:00
parent 3d14d07e4e
commit dacf88c692

View File

@@ -1,11 +1,12 @@
import os import os
from typing import Callable, Optional, Set, Union, cast, Dict from typing import Callable, Dict, Optional, Union, cast
from mypy.checker import TypeChecker from mypy.checker import TypeChecker
from mypy.nodes import MemberExpr, TypeInfo from mypy.nodes import MemberExpr, TypeInfo
from mypy.options import Options from mypy.options import Options
from mypy.plugin import AttributeContext, ClassDefContext, FunctionContext, MethodContext, Plugin from mypy.plugin import AttributeContext, ClassDefContext, FunctionContext, MethodContext, Plugin
from mypy.types import AnyType, Instance, Type, TypeOfAny, TypeType, UnionType from mypy.types import AnyType, Instance, Type, TypeOfAny, TypeType, UnionType
from mypy_django_plugin import helpers, monkeypatch from mypy_django_plugin import helpers, monkeypatch
from mypy_django_plugin.config import Config from mypy_django_plugin.config import Config
from mypy_django_plugin.transformers import fields, init_create from mypy_django_plugin.transformers import fields, init_create
@@ -211,10 +212,6 @@ class DjangoPlugin(Plugin):
def get_function_hook(self, fullname: str def get_function_hook(self, fullname: str
) -> Optional[Callable[[FunctionContext], Type]]: ) -> Optional[Callable[[FunctionContext], Type]]:
sym = self.lookup_fully_qualified(fullname)
if sym and isinstance(sym.node, TypeInfo) and sym.node.has_base(helpers.FIELD_FULLNAME):
return fields.adjust_return_type_of_field_instantiation
if fullname == 'django.contrib.auth.get_user_model': if fullname == 'django.contrib.auth.get_user_model':
return return_user_model_hook return return_user_model_hook
@@ -223,24 +220,34 @@ class DjangoPlugin(Plugin):
return determine_proper_manager_type return determine_proper_manager_type
sym = self.lookup_fully_qualified(fullname) sym = self.lookup_fully_qualified(fullname)
if sym and isinstance(sym.node, TypeInfo): if sym is not None and isinstance(sym.node, TypeInfo):
if sym.node.has_base(helpers.FIELD_FULLNAME):
return fields.adjust_return_type_of_field_instantiation
if sym.node.metadata.get('django', {}).get('generated_init'): if sym.node.metadata.get('django', {}).get('generated_init'):
return init_create.redefine_and_typecheck_model_init return init_create.redefine_and_typecheck_model_init
def get_method_hook(self, fullname: str def get_method_hook(self, fullname: str
) -> Optional[Callable[[MethodContext], Type]]: ) -> Optional[Callable[[MethodContext], Type]]:
if fullname in {'django.apps.registry.Apps.get_model',
'django.db.migrations.state.StateApps.get_model'}:
return determine_model_cls_from_string_for_migrations
manager_classes = self._get_current_manager_bases() manager_classes = self._get_current_manager_bases()
class_fullname, _, method_name = fullname.rpartition('.') class_fullname, _, method_name = fullname.rpartition('.')
if class_fullname in manager_classes and method_name == 'create': if class_fullname in manager_classes and method_name == 'create':
return init_create.redefine_and_typecheck_model_create return init_create.redefine_and_typecheck_model_create
if fullname in {'django.apps.registry.Apps.get_model',
'django.db.migrations.state.StateApps.get_model'}:
return determine_model_cls_from_string_for_migrations
return None return None
def get_base_class_hook(self, fullname: str def get_base_class_hook(self, fullname: str
) -> Optional[Callable[[ClassDefContext], None]]: ) -> Optional[Callable[[ClassDefContext], None]]:
if fullname == helpers.DUMMY_SETTINGS_BASE_CLASS:
settings_modules = ['django.conf.global_settings']
if self.django_settings_module:
settings_modules.append(self.django_settings_module)
return AddSettingValuesToDjangoConfObject(settings_modules,
self.config.ignore_missing_settings)
if fullname in self._get_current_model_bases(): if fullname in self._get_current_model_bases():
return transform_model_class return transform_model_class
@@ -250,17 +257,13 @@ class DjangoPlugin(Plugin):
if fullname in self._get_current_modelform_bases(): if fullname in self._get_current_modelform_bases():
return transform_modelform_class return transform_modelform_class
if fullname == helpers.DUMMY_SETTINGS_BASE_CLASS:
settings_modules = ['django.conf.global_settings']
if self.django_settings_module:
settings_modules.append(self.django_settings_module)
return AddSettingValuesToDjangoConfObject(settings_modules,
self.config.ignore_missing_settings)
return None return None
def get_attribute_hook(self, fullname: str def get_attribute_hook(self, fullname: str
) -> Optional[Callable[[AttributeContext], Type]]: ) -> Optional[Callable[[AttributeContext], Type]]:
if fullname == 'builtins.object.id':
return return_integer_type_for_id_for_non_defined_primary_key_in_models
module, _, name = fullname.rpartition('.') module, _, name = fullname.rpartition('.')
sym = self.lookup_fully_qualified('django.conf.LazySettings') sym = self.lookup_fully_qualified('django.conf.LazySettings')
if sym and isinstance(sym.node, TypeInfo): if sym and isinstance(sym.node, TypeInfo):
@@ -268,9 +271,6 @@ class DjangoPlugin(Plugin):
if module == 'builtins.object' and name in metadata: if module == 'builtins.object' and name in metadata:
return ExtractSettingType(module_fullname=metadata[name]) return ExtractSettingType(module_fullname=metadata[name])
if fullname == 'builtins.object.id':
return return_integer_type_for_id_for_non_defined_primary_key_in_models
return extract_and_return_primary_key_of_bound_related_field_parameter return extract_and_return_primary_key_of_bound_related_field_parameter