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) 297