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,14 +44,8 @@ def _infer_scalar_field(inference_state, field_name, field_tree_instance):
         return attribute.execute_with_values()
 
 
-def _infer_field(cls, field_name):
-    inference_state = cls.inference_state
-    for field_tree_instance in field_name.infer():
-        scalar_field = _infer_scalar_field(inference_state, field_name, field_tree_instance)
-        if scalar_field is not None:
-            return scalar_field
-
-        if field_tree_instance.py__name__() == 'ForeignKey':
+@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()
@@ -63,13 +55,30 @@ def _infer_field(cls, field_name):
                 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()
-                            )
+                    for v in module.py__getattribute__(foreign_key_class_name):
+                        if v.is_class():
+                            yield v
                 elif value.is_class():
-                            return value.execute_with_values()
+                    yield value
+
+
+def _infer_field(cls, field_name):
+    inference_state = cls.inference_state
+    for field_tree_instance in field_name.infer():
+        scalar_field = _infer_scalar_field(inference_state, field_name, field_tree_instance)
+        if scalar_field is not None:
+            return scalar_field
+
+        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]),))

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