Infer many to many fields

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':
+        name = field_tree_instance.py__name__()
-            if isinstance(field_tree_instance, TreeInstance):
+        is_many_to_many = name == 'ManyToManyField'
-                # TODO private access..
+        if name == 'ForeignKey' or is_many_to_many:
-                argument_iterator = field_tree_instance._arguments.unpack()
+            values = _get_foreign_key_values(cls, field_tree_instance)
-                key, lazy_values = next(argument_iterator, (None, None))
+            if is_many_to_many:
-                if key is None and lazy_values is not None:
+                return ValueSet(filter(None, [
-                    for value in lazy_values.infer():
+                    _create_manager_for(v, 'RelatedManager') for v in values
-                        if value.py__name__() == 'str':
+                ]))
-                            foreign_key_class_name = value.get_safe_value()
+            else:
-                            module = cls.get_root_context()
+                return values.execute_with_values()
-                            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()
 
     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]),))

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