From 9de5ab2037a5727a7df873bf0597214f7f3ca58d Mon Sep 17 00:00:00 2001 From: Dave Halter Date: Sat, 13 Jun 2020 20:55:31 +0200 Subject: [PATCH] Make it possible to complete on QuerySet methods, fixes #1587 --- jedi/plugins/django.py | 31 +++++++++++++++++++++++++++++++ test/completion/django.py | 17 +++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/jedi/plugins/django.py b/jedi/plugins/django.py index 6e262a12..098262df 100644 --- a/jedi/plugins/django.py +++ b/jedi/plugins/django.py @@ -13,6 +13,7 @@ from jedi.inference.value.klass import ClassMixin from jedi.inference.gradual.base import GenericClass from jedi.inference.gradual.generics import TupleGenericManager from jedi.inference.signature import AbstractSignature +from jedi.inference.value.function import FunctionMixin mapping = { @@ -171,6 +172,17 @@ def get_metaclass_filters(func): def tree_name_to_values(func): def wrapper(inference_state, context, tree_name): result = func(inference_state, context, tree_name) + if tree_name.value in ('create', 'filter', 'exclude', 'update', 'get', + 'get_or_create', 'update_or_create'): + for v in result: + if v.get_qualified_names() == ('_BaseQuerySet', tree_name.value) \ + and v.parent_context.is_module() \ + and v.parent_context.py__name__() == 'django.db.models.query': + qs = context.get_value() + generics = qs.get_generics() + if len(generics) >= 1: + return ValueSet(QuerySetMethodWrapper(v, model) + for model in generics[0]) if tree_name.value == 'BaseManager' and context.is_module() \ and context.py__name__() == 'django.db.models.manager': return ValueSet(ManagerWrapper(r) for r in result) @@ -257,3 +269,22 @@ class DjangoParamName(BaseTreeParamName): def infer(self): return self._field_name.infer() + + +class QuerySetMethodWrapper(ValueWrapper): + def __init__(self, method, model_cls): + super(QuerySetMethodWrapper, self).__init__(method) + self._model_cls = model_cls + + def py__get__(self, instance, class_value): + return ValueSet({QuerySetBoundMethodWrapper(v, self._model_cls) + for v in self._wrapped_value.py__get__(instance, class_value)}) + + +class QuerySetBoundMethodWrapper(ValueWrapper): + def __init__(self, method, model_cls): + super(QuerySetBoundMethodWrapper, self).__init__(method) + self._model_cls = model_cls + + def get_signatures(self): + return _get_signatures(self._model_cls) diff --git a/test/completion/django.py b/test/completion/django.py index 5387e990..533cab6e 100644 --- a/test/completion/django.py +++ b/test/completion/django.py @@ -273,3 +273,20 @@ Inherited(category_fk) Inherited(attached_o2) #? 18 ['tags_m2m='] Inherited(tags_m2m) + +#? 32 ['tags_m2m='] +Inherited.objects.create(tags_m2) +#? 32 ['tags_m2m='] +Inherited.objects.filter(tags_m2) +#? 35 ['char_field='] +Inherited.objects.exclude(char_fiel) +#? 34 ['char_field='] +Inherited.objects.update(char_fiel) +#? 32 ['email_field='] +Inherited.objects.get(email_fiel) +#? 44 ['category_fk2='] +Inherited.objects.get_or_create(category_fk2) +#? 44 ['uuid_field='] +Inherited.objects.update_or_create(uuid_fiel) +#? 48 ['char_field='] +Inherited.objects.exclude(pk=3).filter(char_fiel)