Django classes plugin

jedi/plugins/django.py

@@ -0,0 +1,103 @@
+"""
+Module is used to infer Django model fields.
+Bugs:
+    - Can't infer User model.
+    - Can't infer ManyToManyField.
+"""
+from jedi.inference.base_value import LazyValueWrapper
+from jedi.inference.utils import safe_property
+from jedi.inference.filters import ParserTreeFilter, DictFilter
+from jedi.inference.names import ValueName
+from jedi.inference.value.instance import TreeInstance
+
+
+def new_dict_filter(cls):
+    filter_ = ParserTreeFilter(parent_context=cls.as_context())
+    res = {f.string_name: _infer_field(cls, f) for f in filter_.values()}
+    return [DictFilter({x: y for x, y in res.items() if y is not None})]
+
+
+class DjangoModelField(LazyValueWrapper):
+    def __init__(self, cls, name):
+        self.inference_state = cls.inference_state
+        self._cls = cls  # Corresponds to super().__self__
+        self._name = name
+        self.tree_node = self._name.tree_name
+
+    @safe_property
+    def name(self):
+        return ValueName(self, self._name.tree_name)
+
+    def _get_wrapped_value(self):
+        obj, = self._cls.execute_with_values()
+        return obj
+
+
+mapping = {
+    'IntegerField': (None, 'int'),
+    'BigIntegerField': (None, 'int'),
+    'PositiveIntegerField': (None, 'int'),
+    'SmallIntegerField': (None, 'int'),
+    'CharField': (None, 'str'),
+    'TextField': (None, 'str'),
+    'EmailField': (None, 'str'),
+    'FloatField': (None, 'float'),
+    'BinaryField': (None, 'bytes'),
+    'BooleanField': (None, 'bool'),
+    'DecimalField': ('decimal', 'Decimal'),
+    'TimeField': ('datetime', 'time'),
+    'DurationField': ('datetime', 'timedelta'),
+    'DateField': ('datetime', 'date'),
+    'DateTimeField': ('datetime', 'datetime'),
+}
+
+
+def _infer_scalar_field(cls, field, field_tree_instance):
+    if field_tree_instance.name.string_name not in mapping:
+        return None
+
+    module_name, attribute_name = mapping[field_tree_instance.name.string_name]
+    if module_name is None:
+        module = cls.inference_state.builtins_module
+    else:
+        module = cls.inference_state.import_module((module_name,))
+
+    attribute, = module.py__getattribute__(attribute_name)
+    return DjangoModelField(attribute, field).name
+
+
+def _infer_field(cls, field):
+    field_tree_instance, = field.infer()
+    scalar_field = _infer_scalar_field(cls, field, field_tree_instance)
+    if scalar_field:
+        return scalar_field
+
+    if field_tree_instance.name.string_name == 'ForeignKey':
+        if isinstance(field_tree_instance, TreeInstance):
+             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.name.string_name == 'str':
+                         foreign_key_class_name = value.get_safe_value()
+                         for v in cls.parent_context.py__getattribute__(foreign_key_class_name):
+                             return DjangoModelField(v, field).name
+                     else:
+                         return DjangoModelField(value, field).name
+
+    print('django plugin: fail to infer `{}` from class `{}`'.format(
+        field.string_name, cls.name.string_name,
+    ))
+
+
+def get_metaclass_filters(func):
+    def wrapper(cls, metaclasses):
+        for metaclass in metaclasses:
+            if metaclass.py__name__() == 'ModelBase' \
+                    and metaclass.get_root_context().py__name__() == 'django.db.models.base':
+                django_dict_filter = new_dict_filter(cls)
+                if django_dict_filter is not None:
+                    return django_dict_filter
+
+        return func(cls, metaclasses)
+    return wrapper

jedi/plugins/registry.py

@@ -5,7 +5,8 @@
 from jedi.plugins import stdlib
 from jedi.plugins import flask
 from jedi.plugins import pytest
+from jedi.plugins import django
 from jedi.plugins import plugin_manager
 
 
-plugin_manager.register(stdlib, flask, pytest)
+plugin_manager.register(stdlib, flask, pytest, django)

setup.py

@@ -43,6 +43,7 @@
               'docopt',
               # coloroma for colored debug output
               'colorama',
+              'Django',
           ],
           'qa': [
               'flake8==3.7.9',

test/completion/django.py

@@ -0,0 +1,73 @@
+import pytest
+import datetime
+import decimal
+
+from django.db import models
+
+
+class Tag(models.Model):
+    tag_name = models.CharField()
+
+
+class Category(models.Model):
+    category_name = models.CharField()
+
+
+class BusinessModel(models.Model):
+    category_fk = models.ForeignKey(Category)
+    integer_field = models.IntegerField()
+    big_integer_field = models.BigIntegerField()
+    positive_integer_field = models.PositiveIntegerField()
+    small_integer_field = models.SmallIntegerField()
+    char_field = models.CharField()
+    text_field = models.TextField()
+    email_field = models.EmailField()
+    float_field = models.FloatField()
+    binary_field = models.BinaryField()
+    boolean_field = models.BooleanField()
+    decimal_field = models.DecimalField()
+    time_field = models.TimeField()
+    duration_field = models.DurationField()
+    date_field = models.DateField()
+    date_time_field = models.DateTimeField()
+    tags_m2m = models.ManyToManyField(Tag)
+
+
+model_instance = BusinessModel()
+#? int()
+model_instance.integer_field
+#? int()
+model_instance.big_integer_field
+#? int()
+model_instance.positive_integer_field
+#? int()
+model_instance.small_integer_field
+#? str()
+model_instance.char_field
+#? str()
+model_instance.text_field
+#? str()
+model_instance.email_field
+#? float()
+model_instance.float_field
+#? bytes()
+model_instance.binary_field
+#? bool()
+model_instance.boolean_field
+#? decimal.Decimal()
+model_instance.decimal_field
+#? datetime.time()
+model_instance.time_field
+#? datetime.timedelta()
+model_instance.duration_field
+#? datetime.date()
+model_instance.date_field
+#? datetime.datetime()
+model_instance.date_time_field
+#? Category()
+model_instance.category_fk
+#? str()
+model_instance.category_fk.category_name
+# TODO: implement many to many field support
+model_instance.tags_m2m
+