Skip to content

Commit

Permalink
optimize hooks a bit
Browse files Browse the repository at this point in the history
  • Loading branch information
mkurnikov committed Feb 21, 2019
1 parent 3d14d07 commit dacf88c
Showing 1 changed file with 20 additions and 20 deletions.
40 changes: 20 additions & 20 deletions mypy_django_plugin/main.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
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.nodes import MemberExpr, TypeInfo
from mypy.options import Options
from mypy.plugin import AttributeContext, ClassDefContext, FunctionContext, MethodContext, Plugin
from mypy.types import AnyType, Instance, Type, TypeOfAny, TypeType, UnionType

from mypy_django_plugin import helpers, monkeypatch
from mypy_django_plugin.config import Config
from mypy_django_plugin.transformers import fields, init_create
Expand Down Expand Up @@ -211,10 +212,6 @@ def _get_current_modelform_bases(self) -> Dict[str, int]:

def get_function_hook(self, fullname: str
) -> 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':
return return_user_model_hook

Expand All @@ -223,24 +220,34 @@ def get_function_hook(self, fullname: str
return determine_proper_manager_type

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'):
return init_create.redefine_and_typecheck_model_init

def get_method_hook(self, fullname: str
) -> 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()
class_fullname, _, method_name = fullname.rpartition('.')
if class_fullname in manager_classes and method_name == '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

def get_base_class_hook(self, fullname: str
) -> 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():
return transform_model_class

Expand All @@ -250,27 +257,20 @@ def get_base_class_hook(self, fullname: str
if fullname in self._get_current_modelform_bases():
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

def get_attribute_hook(self, fullname: str
) -> 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('.')
sym = self.lookup_fully_qualified('django.conf.LazySettings')
if sym and isinstance(sym.node, TypeInfo):
metadata = get_settings_metadata(sym.node)
if module == 'builtins.object' and name in metadata:
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


Expand Down

0 comments on commit dacf88c

Please sign in to comment.