]>
Commit | Line | Data |
---|---|---|
53e6db90 DC |
1 | """ |
2 | Module is used to infer Django model fields. | |
3 | """ | |
4 | from inspect import Parameter | |
5 | ||
6 | from jedi import debug | |
7 | from jedi.inference.cache import inference_state_function_cache | |
8 | from jedi.inference.base_value import ValueSet, iterator_to_value_set, ValueWrapper | |
9 | from jedi.inference.filters import DictFilter, AttributeOverwrite | |
10 | from jedi.inference.names import NameWrapper, BaseTreeParamName | |
11 | from jedi.inference.compiled.value import EmptyCompiledName | |
12 | from jedi.inference.value.instance import TreeInstance | |
13 | from jedi.inference.value.klass import ClassMixin | |
14 | from jedi.inference.gradual.base import GenericClass | |
15 | from jedi.inference.gradual.generics import TupleGenericManager | |
16 | from jedi.inference.signature import AbstractSignature | |
17 | ||
18 | ||
19 | mapping = { | |
20 | 'IntegerField': (None, 'int'), | |
21 | 'BigIntegerField': (None, 'int'), | |
22 | 'PositiveIntegerField': (None, 'int'), | |
23 | 'SmallIntegerField': (None, 'int'), | |
24 | 'CharField': (None, 'str'), | |
25 | 'TextField': (None, 'str'), | |
26 | 'EmailField': (None, 'str'), | |
27 | 'GenericIPAddressField': (None, 'str'), | |
28 | 'URLField': (None, 'str'), | |
29 | 'FloatField': (None, 'float'), | |
30 | 'BinaryField': (None, 'bytes'), | |
31 | 'BooleanField': (None, 'bool'), | |
32 | 'DecimalField': ('decimal', 'Decimal'), | |
33 | 'TimeField': ('datetime', 'time'), | |
34 | 'DurationField': ('datetime', 'timedelta'), | |
35 | 'DateField': ('datetime', 'date'), | |
36 | 'DateTimeField': ('datetime', 'datetime'), | |
37 | 'UUIDField': ('uuid', 'UUID'), | |
38 | } | |
39 | ||
40 | _FILTER_LIKE_METHODS = ('create', 'filter', 'exclude', 'update', 'get', | |
41 | 'get_or_create', 'update_or_create') | |
42 | ||
43 | ||
44 | @inference_state_function_cache() | |
45 | def _get_deferred_attributes(inference_state): | |
46 | return inference_state.import_module( | |
47 | ('django', 'db', 'models', 'query_utils') | |
48 | ).py__getattribute__('DeferredAttribute').execute_annotation() | |
49 | ||
50 | ||
51 | def _infer_scalar_field(inference_state, field_name, field_tree_instance, is_instance): | |
52 | try: | |
53 | module_name, attribute_name = mapping[field_tree_instance.py__name__()] | |
54 | except KeyError: | |
55 | return None | |
56 | ||
57 | if not is_instance: | |
58 | return _get_deferred_attributes(inference_state) | |
59 | ||
60 | if module_name is None: | |
61 | module = inference_state.builtins_module | |
62 | else: | |
63 | module = inference_state.import_module((module_name,)) | |
64 | ||
65 | for attribute in module.py__getattribute__(attribute_name): | |
66 | return attribute.execute_with_values() | |
67 | ||
68 | ||
69 | @iterator_to_value_set | |
70 | def _get_foreign_key_values(cls, field_tree_instance): | |
71 | if isinstance(field_tree_instance, TreeInstance): | |
72 | # TODO private access.. | |
73 | argument_iterator = field_tree_instance._arguments.unpack() | |
74 | key, lazy_values = next(argument_iterator, (None, None)) | |
75 | if key is None and lazy_values is not None: | |
76 | for value in lazy_values.infer(): | |
77 | if value.py__name__() == 'str': | |
78 | foreign_key_class_name = value.get_safe_value() | |
79 | module = cls.get_root_context() | |
80 | for v in module.py__getattribute__(foreign_key_class_name): | |
81 | if v.is_class(): | |
82 | yield v | |
83 | elif value.is_class(): | |
84 | yield value | |
85 | ||
86 | ||
87 | def _infer_field(cls, field_name, is_instance): | |
88 | inference_state = cls.inference_state | |
89 | result = field_name.infer() | |
90 | for field_tree_instance in result: | |
91 | scalar_field = _infer_scalar_field( | |
92 | inference_state, field_name, field_tree_instance, is_instance) | |
93 | if scalar_field is not None: | |
94 | return scalar_field | |
95 | ||
96 | name = field_tree_instance.py__name__() | |
97 | is_many_to_many = name == 'ManyToManyField' | |
98 | if name in ('ForeignKey', 'OneToOneField') or is_many_to_many: | |
99 | if not is_instance: | |
100 | return _get_deferred_attributes(inference_state) | |
101 | ||
102 | values = _get_foreign_key_values(cls, field_tree_instance) | |
103 | if is_many_to_many: | |
104 | return ValueSet(filter(None, [ | |
105 | _create_manager_for(v, 'RelatedManager') for v in values | |
106 | ])) | |
107 | else: | |
108 | return values.execute_with_values() | |
109 | ||
110 | debug.dbg('django plugin: fail to infer `%s` from class `%s`', | |
111 | field_name.string_name, cls.py__name__()) | |
112 | return result | |
113 | ||
114 | ||
115 | class DjangoModelName(NameWrapper): | |
116 | def __init__(self, cls, name, is_instance): | |
117 | super().__init__(name) | |
118 | self._cls = cls | |
119 | self._is_instance = is_instance | |
120 | ||
121 | def infer(self): | |
122 | return _infer_field(self._cls, self._wrapped_name, self._is_instance) | |
123 | ||
124 | ||
125 | def _create_manager_for(cls, manager_cls='BaseManager'): | |
126 | managers = cls.inference_state.import_module( | |
127 | ('django', 'db', 'models', 'manager') | |
128 | ).py__getattribute__(manager_cls) | |
129 | for m in managers: | |
130 | if m.is_class_mixin(): | |
131 | generics_manager = TupleGenericManager((ValueSet([cls]),)) | |
132 | for c in GenericClass(m, generics_manager).execute_annotation(): | |
133 | return c | |
134 | return None | |
135 | ||
136 | ||
137 | def _new_dict_filter(cls, is_instance): | |
138 | filters = list(cls.get_filters( | |
139 | is_instance=is_instance, | |
140 | include_metaclasses=False, | |
141 | include_type_when_class=False) | |
142 | ) | |
143 | dct = { | |
144 | name.string_name: DjangoModelName(cls, name, is_instance) | |
145 | for filter_ in reversed(filters) | |
146 | for name in filter_.values() | |
147 | } | |
148 | if is_instance: | |
149 | # Replace the objects with a name that amounts to nothing when accessed | |
150 | # in an instance. This is not perfect and still completes "objects" in | |
151 | # that case, but it at least not inferes stuff like `.objects.filter`. | |
152 | # It would be nicer to do that in a better way, so that it also doesn't | |
153 | # show up in completions, but it's probably just not worth doing that | |
154 | # for the extra amount of work. | |
155 | dct['objects'] = EmptyCompiledName(cls.inference_state, 'objects') | |
156 | ||
157 | return DictFilter(dct) | |
158 | ||
159 | ||
160 | def is_django_model_base(value): | |
161 | return value.py__name__() == 'ModelBase' \ | |
162 | and value.get_root_context().py__name__() == 'django.db.models.base' | |
163 | ||
164 | ||
165 | def get_metaclass_filters(func): | |
166 | def wrapper(cls, metaclasses, is_instance): | |
167 | for metaclass in metaclasses: | |
168 | if is_django_model_base(metaclass): | |
169 | return [_new_dict_filter(cls, is_instance)] | |
170 | ||
171 | return func(cls, metaclasses, is_instance) | |
172 | return wrapper | |
173 | ||
174 | ||
175 | def tree_name_to_values(func): | |
176 | def wrapper(inference_state, context, tree_name): | |
177 | result = func(inference_state, context, tree_name) | |
178 | if tree_name.value in _FILTER_LIKE_METHODS: | |
179 | # Here we try to overwrite stuff like User.objects.filter. We need | |
180 | # this to make sure that keyword param completion works on these | |
181 | # kind of methods. | |
182 | for v in result: | |
183 | if v.get_qualified_names() == ('_BaseQuerySet', tree_name.value) \ | |
184 | and v.parent_context.is_module() \ | |
185 | and v.parent_context.py__name__() == 'django.db.models.query': | |
186 | qs = context.get_value() | |
187 | generics = qs.get_generics() | |
188 | if len(generics) >= 1: | |
189 | return ValueSet(QuerySetMethodWrapper(v, model) | |
190 | for model in generics[0]) | |
191 | ||
192 | elif tree_name.value == 'BaseManager' and context.is_module() \ | |
193 | and context.py__name__() == 'django.db.models.manager': | |
194 | return ValueSet(ManagerWrapper(r) for r in result) | |
195 | ||
196 | elif tree_name.value == 'Field' and context.is_module() \ | |
197 | and context.py__name__() == 'django.db.models.fields': | |
198 | return ValueSet(FieldWrapper(r) for r in result) | |
199 | return result | |
200 | return wrapper | |
201 | ||
202 | ||
203 | def _find_fields(cls): | |
204 | for name in _new_dict_filter(cls, is_instance=False).values(): | |
205 | for value in name.infer(): | |
206 | if value.name.get_qualified_names(include_module_names=True) \ | |
207 | == ('django', 'db', 'models', 'query_utils', 'DeferredAttribute'): | |
208 | yield name | |
209 | ||
210 | ||
211 | def _get_signatures(cls): | |
212 | return [DjangoModelSignature(cls, field_names=list(_find_fields(cls)))] | |
213 | ||
214 | ||
215 | def get_metaclass_signatures(func): | |
216 | def wrapper(cls, metaclasses): | |
217 | for metaclass in metaclasses: | |
218 | if is_django_model_base(metaclass): | |
219 | return _get_signatures(cls) | |
220 | return func(cls, metaclass) | |
221 | return wrapper | |
222 | ||
223 | ||
224 | class ManagerWrapper(ValueWrapper): | |
225 | def py__getitem__(self, index_value_set, contextualized_node): | |
226 | return ValueSet( | |
227 | GenericManagerWrapper(generic) | |
228 | for generic in self._wrapped_value.py__getitem__( | |
229 | index_value_set, contextualized_node) | |
230 | ) | |
231 | ||
232 | ||
233 | class GenericManagerWrapper(AttributeOverwrite, ClassMixin): | |
234 | def py__get__on_class(self, calling_instance, instance, class_value): | |
235 | return calling_instance.class_value.with_generics( | |
236 | (ValueSet({class_value}),) | |
237 | ).py__call__(calling_instance._arguments) | |
238 | ||
239 | def with_generics(self, generics_tuple): | |
240 | return self._wrapped_value.with_generics(generics_tuple) | |
241 | ||
242 | ||
243 | class FieldWrapper(ValueWrapper): | |
244 | def py__getitem__(self, index_value_set, contextualized_node): | |
245 | return ValueSet( | |
246 | GenericFieldWrapper(generic) | |
247 | for generic in self._wrapped_value.py__getitem__( | |
248 | index_value_set, contextualized_node) | |
249 | ) | |
250 | ||
251 | ||
252 | class GenericFieldWrapper(AttributeOverwrite, ClassMixin): | |
253 | def py__get__on_class(self, calling_instance, instance, class_value): | |
254 | # This is mostly an optimization to avoid Jedi aborting inference, | |
255 | # because of too many function executions of Field.__get__. | |
256 | return ValueSet({calling_instance}) | |
257 | ||
258 | ||
259 | class DjangoModelSignature(AbstractSignature): | |
260 | def __init__(self, value, field_names): | |
261 | super().__init__(value) | |
262 | self._field_names = field_names | |
263 | ||
264 | def get_param_names(self, resolve_stars=False): | |
265 | return [DjangoParamName(name) for name in self._field_names] | |
266 | ||
267 | ||
268 | class DjangoParamName(BaseTreeParamName): | |
269 | def __init__(self, field_name): | |
270 | super().__init__(field_name.parent_context, field_name.tree_name) | |
271 | self._field_name = field_name | |
272 | ||
273 | def get_kind(self): | |
274 | return Parameter.KEYWORD_ONLY | |
275 | ||
276 | def infer(self): | |
277 | return self._field_name.infer() | |
278 | ||
279 | ||
280 | class QuerySetMethodWrapper(ValueWrapper): | |
281 | def __init__(self, method, model_cls): | |
282 | super().__init__(method) | |
283 | self._model_cls = model_cls | |
284 | ||
285 | def py__get__(self, instance, class_value): | |
286 | return ValueSet({QuerySetBoundMethodWrapper(v, self._model_cls) | |
287 | for v in self._wrapped_value.py__get__(instance, class_value)}) | |
288 | ||
289 | ||
290 | class QuerySetBoundMethodWrapper(ValueWrapper): | |
291 | def __init__(self, method, model_cls): | |
292 | super().__init__(method) | |
293 | self._model_cls = model_cls | |
294 | ||
295 | def get_signatures(self): | |
296 | return _get_signatures(self._model_cls) |