diff --git a/jedi/plugins/django.py b/jedi/plugins/django.py index 3c0b3cbe..0400c6a9 100644 --- a/jedi/plugins/django.py +++ b/jedi/plugins/django.py @@ -1,10 +1,8 @@ """ Module is used to infer Django model fields. -Bugs: - - Can't infer ManyToManyField. """ from jedi import debug -from jedi.inference.base_value import ValueSet +from jedi.inference.base_value import ValueSet, iterator_to_value_set from jedi.inference.filters import ParserTreeFilter, DictFilter from jedi.inference.names import NameWrapper from jedi.inference.value.instance import TreeInstance @@ -46,6 +44,24 @@ def _infer_scalar_field(inference_state, field_name, field_tree_instance): return attribute.execute_with_values() +@iterator_to_value_set +def _get_foreign_key_values(cls, field_tree_instance): + if isinstance(field_tree_instance, TreeInstance): + # TODO private access.. + argument_iterator = field_tree_instance._arguments.unpack() + key, lazy_values = next(argument_iterator, (None, None)) + if key is None and lazy_values is not None: + for value in lazy_values.infer(): + if value.py__name__() == 'str': + foreign_key_class_name = value.get_safe_value() + module = cls.get_root_context() + for v in module.py__getattribute__(foreign_key_class_name): + if v.is_class(): + yield v + elif value.is_class(): + yield value + + def _infer_field(cls, field_name): inference_state = cls.inference_state for field_tree_instance in field_name.infer(): @@ -53,23 +69,16 @@ def _infer_field(cls, field_name): if scalar_field is not None: return scalar_field - if field_tree_instance.py__name__() == 'ForeignKey': - if isinstance(field_tree_instance, TreeInstance): - # TODO private access.. - argument_iterator = field_tree_instance._arguments.unpack() - key, lazy_values = next(argument_iterator, (None, None)) - if key is None and lazy_values is not None: - for value in lazy_values.infer(): - if value.py__name__() == 'str': - foreign_key_class_name = value.get_safe_value() - module = cls.get_root_context() - return ValueSet.from_sets( - v.execute_with_values() - for v in module.py__getattribute__(foreign_key_class_name) - if v.is_class() - ) - elif value.is_class(): - return value.execute_with_values() + name = field_tree_instance.py__name__() + is_many_to_many = name == 'ManyToManyField' + if name == 'ForeignKey' or is_many_to_many: + values = _get_foreign_key_values(cls, field_tree_instance) + if is_many_to_many: + return ValueSet(filter(None, [ + _create_manager_for(v, 'RelatedManager') for v in values + ])) + else: + return values.execute_with_values() debug.dbg('django plugin: fail to infer `%s` from class `%s`', field_name.string_name, cls.py__name__()) @@ -85,10 +94,10 @@ class DjangoModelName(NameWrapper): return _infer_field(self._cls, self._wrapped_name) -def _create_manager_for(cls): +def _create_manager_for(cls, manager_cls='BaseManager'): managers = cls.inference_state.import_module( ('django', 'db', 'models', 'manager') - ).py__getattribute__('BaseManager') + ).py__getattribute__(manager_cls) for m in managers: if m.is_class() and not m.is_compiled(): generics_manager = TupleGenericManager((ValueSet([cls]),)) diff --git a/test/completion/django.py b/test/completion/django.py index d0349601..4c2cf6ec 100644 --- a/test/completion/django.py +++ b/test/completion/django.py @@ -91,15 +91,19 @@ model_instance.category_fk.category_name model_instance.category_fk2 #? str() model_instance.category_fk2.category_name -#? models.ForeignKey() +#? model_instance.category_fk3 #? model_instance.category_fk4 -#? models.ForeignKey() +#? model_instance.category_fk5 -#? models.ManyToManyField() +#? models.manager.RelatedManager() model_instance.tags_m2m +#? Tag() +model_instance.tags_m2m.get() +#? ['add'] +model_instance.tags_m2m.add #? model_instance.unidentifiable