def extract_proper_type_queryset_values(ctx: MethodContext, django_context: DjangoContext) -> MypyType: # called on QuerySet, return QuerySet of something assert isinstance(ctx.type, Instance) assert isinstance(ctx.default_return_type, Instance) model_type = _extract_model_type_from_queryset(ctx.type) if model_type is None: return AnyType(TypeOfAny.from_omitted_generics) model_cls = django_context.get_model_class_by_fullname(model_type.type.fullname) if model_cls is None: return ctx.default_return_type field_lookups = resolve_field_lookups(ctx.args[0], django_context) if field_lookups is None: return AnyType(TypeOfAny.from_error) if len(field_lookups) == 0: for field in django_context.get_model_fields(model_cls): field_lookups.append(field.attname) column_types: 'OrderedDict[str, MypyType]' = OrderedDict() for field_lookup in field_lookups: field_lookup_type = get_field_type_from_lookup(ctx, django_context, model_cls, lookup=field_lookup, method='values') if field_lookup_type is None: return helpers.reparametrize_instance(ctx.default_return_type, [model_type, AnyType(TypeOfAny.from_error)]) column_types[field_lookup] = field_lookup_type row_type = helpers.make_typeddict(ctx.api, column_types, set(column_types.keys())) return helpers.reparametrize_instance(ctx.default_return_type, [model_type, row_type])
def reparametrize_related_field_type(related_field_type: Instance, set_type, get_type) -> Instance: args = [ helpers.convert_any_to_type(related_field_type.args[0], set_type), helpers.convert_any_to_type(related_field_type.args[1], get_type), ] return helpers.reparametrize_instance(related_field_type, new_args=args)
def create_new_model_parametrized_manager(self, name: str, base_manager_info: TypeInfo) -> Instance: bases = [] for original_base in base_manager_info.bases: if self.is_any_parametrized_manager(original_base): if original_base.type is None: raise helpers.IncompleteDefnException() original_base = helpers.reparametrize_instance(original_base, [Instance(self.model_classdef.info, [])]) bases.append(original_base) new_manager_info = self.add_new_class_for_current_module(name, bases) # copy fields to a new manager new_cls_def_context = ClassDefContext(cls=new_manager_info.defn, reason=self.ctx.reason, api=self.api) custom_manager_type = Instance(new_manager_info, [Instance(self.model_classdef.info, [])]) for name, sym in base_manager_info.names.items(): # replace self type with new class, if copying method if isinstance(sym.node, FuncDef): helpers.copy_method_to_another_class( new_cls_def_context, self_type=custom_manager_type, new_method_name=name, method_node=sym.node ) continue new_sym = sym.copy() if isinstance(new_sym.node, Var): new_var = Var(name, type=sym.type) new_var.info = new_manager_info new_var._fullname = new_manager_info.fullname + "." + name new_sym.node = new_var new_manager_info.names[name] = new_sym return custom_manager_type
def run_with_model_cls(self, model_cls: Type[Model]) -> None: for manager_name, manager in model_cls._meta.managers_map.items(): manager_fullname = helpers.get_class_fullname(manager.__class__) manager_info = self.lookup_typeinfo_or_incomplete_defn_error(manager_fullname) if manager_name not in self.model_classdef.info.names: manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])]) self.add_new_node_to_model_class(manager_name, manager_type) else: # creates new MODELNAME_MANAGERCLASSNAME class that represents manager parametrized with current model has_manager_any_base = any(self._is_manager_any(base) for base in manager_info.bases) if has_manager_any_base: custom_model_manager_name = manager.model.__name__ + '_' + manager.__class__.__name__ bases = [] for original_base in manager_info.bases: if self._is_manager_any(original_base): if original_base.type is None: raise helpers.IncompleteDefnException() original_base = helpers.reparametrize_instance(original_base, [Instance(self.model_classdef.info, [])]) bases.append(original_base) current_module = self.api.modules[self.model_classdef.info.module_name] custom_manager_info = helpers.add_new_class_for_module(current_module, custom_model_manager_name, bases=bases, fields=OrderedDict()) custom_manager_type = Instance(custom_manager_info, [Instance(self.model_classdef.info, [])]) self.add_new_node_to_model_class(manager_name, custom_manager_type)
def extract_proper_type_queryset_values_list( ctx: MethodContext, django_context: DjangoContext) -> MypyType: # called on the Instance, returns QuerySet of something assert isinstance(ctx.type, Instance) default_return_type = get_proper_type(ctx.default_return_type) assert isinstance(default_return_type, Instance) model_type = _extract_model_type_from_queryset(ctx.type) if model_type is None: return AnyType(TypeOfAny.from_omitted_generics) model_cls = django_context.get_model_class_by_fullname( model_type.type.fullname) if model_cls is None: return default_return_type flat_expr = helpers.get_call_argument_by_name(ctx, "flat") if flat_expr is not None and isinstance(flat_expr, NameExpr): flat = helpers.parse_bool(flat_expr) else: flat = False named_expr = helpers.get_call_argument_by_name(ctx, "named") if named_expr is not None and isinstance(named_expr, NameExpr): named = helpers.parse_bool(named_expr) else: named = False if flat and named: ctx.api.fail("'flat' and 'named' can't be used together", ctx.context) return helpers.reparametrize_instance( default_return_type, [model_type, AnyType(TypeOfAny.from_error)]) # account for possible None flat = flat or False named = named or False is_annotated = is_annotated_model_fullname(model_type.type.fullname) row_type = get_values_list_row_type(ctx, django_context, model_cls, is_annotated=is_annotated, flat=flat, named=named) return helpers.reparametrize_instance(default_return_type, [model_type, row_type])
def set_descriptor_types_for_field(ctx: FunctionContext) -> Instance: default_return_type = cast(Instance, ctx.default_return_type) is_nullable = False null_expr = helpers.get_call_argument_by_name(ctx, 'null') if null_expr is not None: is_nullable = helpers.parse_bool(null_expr) or False set_type, get_type = get_field_descriptor_types(default_return_type.type, is_nullable) return helpers.reparametrize_instance(default_return_type, [set_type, get_type])
def determine_proper_manager_type(ctx: FunctionContext) -> MypyType: default_return_type = ctx.default_return_type assert isinstance(default_return_type, Instance) outer_model_info = helpers.get_typechecker_api(ctx).scope.active_class() if (outer_model_info is None or not outer_model_info.has_base(fullnames.MODEL_CLASS_FULLNAME)): return default_return_type return helpers.reparametrize_instance(default_return_type, [Instance(outer_model_info, [])])
def run_with_model_cls(self, model_cls: Type[Model]) -> None: for manager_name, manager in model_cls._meta.managers_map.items(): manager_fullname = helpers.get_class_fullname(manager.__class__) manager_info = self.lookup_typeinfo_or_incomplete_defn_error(manager_fullname) if manager_name not in self.model_classdef.info.names: manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])]) self.add_new_node_to_model_class(manager_name, manager_type) else: # create new MODELNAME_MANAGERCLASSNAME class that represents manager parametrized with current model has_manager_any_base = any(self._is_manager_any(base) for base in manager_info.bases) if has_manager_any_base: custom_model_manager_name = manager.model.__name__ + '_' + manager.__class__.__name__ bases = [] for original_base in manager_info.bases: if self._is_manager_any(original_base): if original_base.type is None: if not self.api.final_iteration: self.api.defer() original_base = helpers.reparametrize_instance(original_base, [Instance(self.model_classdef.info, [])]) bases.append(original_base) current_module = self.api.modules[self.model_classdef.info.module_name] custom_manager_info = helpers.add_new_class_for_module(current_module, custom_model_manager_name, bases=bases, fields=OrderedDict()) custom_manager_type = Instance(custom_manager_info, [Instance(self.model_classdef.info, [])]) self.add_new_node_to_model_class(manager_name, custom_manager_type) # add _default_manager if '_default_manager' not in self.model_classdef.info.names: default_manager_fullname = helpers.get_class_fullname(model_cls._meta.default_manager.__class__) default_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(default_manager_fullname) default_manager = Instance(default_manager_info, [Instance(self.model_classdef.info, [])]) self.add_new_node_to_model_class('_default_manager', default_manager) # add related managers for relation in self.django_context.get_model_relations(model_cls): attname = relation.get_accessor_name() if attname is None: # no reverse accessor continue related_model_info = self.lookup_class_typeinfo_or_incomplete_defn_error(relation.related_model) if isinstance(relation, OneToOneRel): self.add_new_node_to_model_class(attname, Instance(related_model_info, [])) continue if isinstance(relation, (ManyToOneRel, ManyToManyRel)): manager_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.RELATED_MANAGER_CLASS_FULLNAME) self.add_new_node_to_model_class(attname, Instance(manager_info, [Instance(related_model_info, [])])) continue
def determine_type_of_array_field(ctx: FunctionContext, django_context: DjangoContext) -> MypyType: default_return_type = set_descriptor_types_for_field(ctx) base_field_arg_type = helpers.get_call_argument_type_by_name( ctx, "base_field") if not base_field_arg_type or not isinstance(base_field_arg_type, Instance): return default_return_type def drop_combinable(_type: MypyType) -> Optional[MypyType]: if isinstance(_type, Instance) and _type.type.has_base( fullnames.COMBINABLE_EXPRESSION_FULLNAME): return None elif isinstance(_type, UnionType): items_without_combinable = [] for item in _type.items: reduced = drop_combinable(item) if reduced is not None: items_without_combinable.append(reduced) if len(items_without_combinable) > 1: return UnionType( items_without_combinable, line=_type.line, column=_type.column, is_evaluated=_type.is_evaluated, uses_pep604_syntax=_type.uses_pep604_syntax, ) elif len(items_without_combinable) == 1: return items_without_combinable[0] else: return None return _type # Both base_field and return type should derive from Field and thus expect 2 arguments assert len(base_field_arg_type.args) == len(default_return_type.args) == 2 args = [] for new_type, default_arg in zip(base_field_arg_type.args, default_return_type.args): # Drop any base_field Combinable type reduced = drop_combinable(new_type) if reduced is None: ctx.api.fail( f"Can't have ArrayField expecting {fullnames.COMBINABLE_EXPRESSION_FULLNAME!r} as data type", ctx.context, ) else: new_type = reduced args.append(helpers.convert_any_to_type(default_arg, new_type)) return helpers.reparametrize_instance(default_return_type, args)
def determine_type_of_array_field(ctx: FunctionContext, django_context: DjangoContext) -> MypyType: default_return_type = set_descriptor_types_for_field(ctx) base_field_arg_type = helpers.get_call_argument_type_by_name(ctx, 'base_field') if not base_field_arg_type or not isinstance(base_field_arg_type, Instance): return default_return_type base_type = base_field_arg_type.args[1] # extract __get__ type args = [] for default_arg in default_return_type.args: args.append(helpers.convert_any_to_type(default_arg, base_type)) return helpers.reparametrize_instance(default_return_type, args)
def fill_descriptor_types_for_related_field( ctx: FunctionContext, django_context: DjangoContext) -> MypyType: current_field = _get_current_field_from_assignment(ctx, django_context) if current_field is None: return AnyType(TypeOfAny.from_error) assert isinstance(current_field, RelatedField) related_model_cls = django_context.fields_context.get_related_model_cls( current_field) related_model = related_model_cls related_model_to_set = related_model_cls if related_model_to_set._meta.proxy_for_model is not None: related_model_to_set = related_model_to_set._meta.proxy_for_model typechecker_api = helpers.get_typechecker_api(ctx) related_model_info = helpers.lookup_class_typeinfo(typechecker_api, related_model) if related_model_info is None: # maybe no type stub related_model_type = AnyType(TypeOfAny.unannotated) else: related_model_type = Instance(related_model_info, []) # type: ignore related_model_to_set_info = helpers.lookup_class_typeinfo( typechecker_api, related_model_to_set) if related_model_to_set_info is None: # maybe no type stub related_model_to_set_type = AnyType(TypeOfAny.unannotated) else: related_model_to_set_type = Instance(related_model_to_set_info, []) # type: ignore default_related_field_type = set_descriptor_types_for_field(ctx) # replace Any with referred_to_type args = [ helpers.convert_any_to_type(default_related_field_type.args[0], related_model_to_set_type), helpers.convert_any_to_type(default_related_field_type.args[1], related_model_type), ] return helpers.reparametrize_instance(default_related_field_type, new_args=args)
def set_descriptor_types_for_field(ctx: FunctionContext, *, is_set_nullable: bool = False, is_get_nullable: bool = False) -> Instance: default_return_type = cast(Instance, ctx.default_return_type) is_nullable = False null_expr = helpers.get_call_argument_by_name(ctx, "null") if null_expr is not None: is_nullable = helpers.parse_bool(null_expr) or False # Allow setting field value to `None` when a field is primary key and has a default that can produce a value default_expr = helpers.get_call_argument_by_name(ctx, "default") primary_key_expr = helpers.get_call_argument_by_name(ctx, "primary_key") if default_expr is not None and primary_key_expr is not None: is_set_nullable = helpers.parse_bool(primary_key_expr) or False set_type, get_type = get_field_descriptor_types( default_return_type.type, is_set_nullable=is_set_nullable or is_nullable, is_get_nullable=is_get_nullable or is_nullable, ) return helpers.reparametrize_instance(default_return_type, [set_type, get_type])
def extract_proper_type_queryset_annotate( ctx: MethodContext, django_context: DjangoContext) -> MypyType: # called on the Instance, returns QuerySet of something assert isinstance(ctx.type, Instance) default_return_type = get_proper_type(ctx.default_return_type) assert isinstance(default_return_type, Instance) model_type = _extract_model_type_from_queryset(ctx.type) if model_type is None: return AnyType(TypeOfAny.from_omitted_generics) api = ctx.api field_types = model_type.type.metadata.get("annotated_field_types") kwargs = gather_kwargs(ctx) if kwargs: # For now, we don't try to resolve the output_field of the field would be, but use Any. added_field_types = { name: AnyType(TypeOfAny.implementation_artifact) for name, typ in kwargs.items() } if field_types is not None: # Annotate was called more than once, so add/update existing field types field_types.update(added_field_types) else: field_types = added_field_types fields_dict = None if field_types is not None: fields_dict = helpers.make_typeddict(api, fields=OrderedDict(field_types), required_keys=set( field_types.keys())) annotated_type = get_or_create_annotated_type(api, model_type, fields_dict=fields_dict) row_type: MypyType if len(default_return_type.args) > 1: original_row_type: MypyType = default_return_type.args[1] row_type = original_row_type if isinstance(original_row_type, TypedDictType): row_type = api.named_generic_type("builtins.dict", [ api.named_generic_type("builtins.str", []), AnyType(TypeOfAny.from_omitted_generics) ]) elif isinstance(original_row_type, TupleType): fallback: Instance = original_row_type.partial_fallback if fallback is not None and fallback.type.has_base( "typing.NamedTuple"): # TODO: Use a NamedTuple which contains the known fields, but also # falls back to allowing any attribute access. row_type = AnyType(TypeOfAny.implementation_artifact) else: row_type = api.named_generic_type( "builtins.tuple", [AnyType(TypeOfAny.from_omitted_generics)]) elif isinstance(original_row_type, Instance) and original_row_type.type.has_base( fullnames.MODEL_CLASS_FULLNAME): row_type = annotated_type else: row_type = annotated_type return helpers.reparametrize_instance(default_return_type, [annotated_type, row_type])