def _get_current_field_from_assignment( ctx: FunctionContext, django_context: DjangoContext) -> Optional[Field]: outer_model_info = helpers.get_typechecker_api(ctx).scope.active_class() if (outer_model_info is None or not helpers.is_model_subclass_info( outer_model_info, django_context)): return None field_name = None for stmt in outer_model_info.defn.defs.body: if isinstance(stmt, AssignmentStmt): if stmt.rvalue == ctx.context: if not isinstance(stmt.lvalues[0], NameExpr): return None field_name = stmt.lvalues[0].name break if field_name is None: return None model_cls = django_context.get_model_class_by_fullname( outer_model_info.fullname) if model_cls is None: return None current_field = model_cls._meta.get_field(field_name) return current_field
def return_proper_field_type_from_get_field( ctx: MethodContext, django_context: DjangoContext) -> MypyType: # Options instance assert isinstance(ctx.type, Instance) model_type = ctx.type.args[0] if not isinstance(model_type, Instance): return _get_field_instance(ctx, fullnames.FIELD_FULLNAME) model_cls = django_context.get_model_class_by_fullname( model_type.type.fullname()) if model_cls is None: return _get_field_instance(ctx, fullnames.FIELD_FULLNAME) field_name_expr = helpers.get_call_argument_by_name(ctx, 'field_name') if field_name_expr is None: return _get_field_instance(ctx, fullnames.FIELD_FULLNAME) field_name = helpers.resolve_string_attribute_value( field_name_expr, ctx, django_context) if field_name is None: return _get_field_instance(ctx, fullnames.FIELD_FULLNAME) try: field = model_cls._meta.get_field(field_name) except FieldDoesNotExist as exc: ctx.api.fail(exc.args[0], ctx.context) return AnyType(TypeOfAny.from_error) field_fullname = helpers.get_class_fullname(field.__class__) return _get_field_instance(ctx, field_fullname)
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 redefine_and_typecheck_model_init( ctx: FunctionContext, django_context: DjangoContext) -> MypyType: assert isinstance(ctx.default_return_type, Instance) model_fullname = ctx.default_return_type.type.fullname model_cls = django_context.get_model_class_by_fullname(model_fullname) if model_cls is None: return ctx.default_return_type return typecheck_model_method(ctx, django_context, model_cls, '__init__')
def redefine_and_typecheck_model_create( ctx: MethodContext, django_context: DjangoContext) -> MypyType: if not isinstance(ctx.default_return_type, Instance): # only work with ctx.default_return_type = model Instance return ctx.default_return_type model_fullname = ctx.default_return_type.type.fullname model_cls = django_context.get_model_class_by_fullname(model_fullname) if model_cls is None: return ctx.default_return_type return typecheck_model_method(ctx, django_context, model_cls, 'create')
def typecheck_queryset_filter(ctx: MethodContext, django_context: DjangoContext) -> MypyType: # Expected formal arguments for filter methods are `*args` and `**kwargs`. We'll only typecheck # `**kwargs`, which means that `arg_names[1]` is what we're interested in. lookup_kwargs = ctx.arg_names[1] provided_lookup_types = ctx.arg_types[1] assert isinstance(ctx.type, Instance) if not ctx.type.args or not isinstance(ctx.type.args[0], Instance): return ctx.default_return_type model_cls_fullname = ctx.type.args[0].type.fullname model_cls = django_context.get_model_class_by_fullname(model_cls_fullname) if model_cls is None: return ctx.default_return_type for lookup_kwarg, provided_type in zip(lookup_kwargs, provided_lookup_types): if lookup_kwarg is None: continue if isinstance(provided_type, Instance) and provided_type.type.has_base( fullnames.COMBINABLE_EXPRESSION_FULLNAME): provided_type = resolve_combinable_type(provided_type, django_context) lookup_type: MypyType if is_annotated_model_fullname(model_cls_fullname): lookup_type = AnyType(TypeOfAny.implementation_artifact) else: lookup_type = django_context.resolve_lookup_expected_type( ctx, model_cls, lookup_kwarg) # Managers as provided_type is not supported yet if isinstance(provided_type, Instance) and helpers.has_any_of_bases( provided_type.type, (fullnames.MANAGER_CLASS_FULLNAME, fullnames.QUERYSET_CLASS_FULLNAME)): return ctx.default_return_type helpers.check_types_compatible( ctx, expected_type=lookup_type, actual_type=provided_type, error_message=f"Incompatible type for lookup {lookup_kwarg!r}:", ) return ctx.default_return_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 typecheck_queryset_filter(ctx: MethodContext, django_context: DjangoContext) -> MypyType: lookup_kwargs = ctx.arg_names[1] provided_lookup_types = ctx.arg_types[1] assert isinstance(ctx.type, Instance) if not ctx.type.args or not isinstance(ctx.type.args[0], Instance): return ctx.default_return_type model_cls_fullname = ctx.type.args[0].type.fullname model_cls = django_context.get_model_class_by_fullname(model_cls_fullname) if model_cls is None: return ctx.default_return_type for lookup_kwarg, provided_type in zip(lookup_kwargs, provided_lookup_types): if lookup_kwarg is None: continue if isinstance(provided_type, Instance) and provided_type.type.has_base( "django.db.models.expressions.Combinable"): provided_type = resolve_combinable_type(provided_type, django_context) lookup_type: MypyType if is_annotated_model_fullname(model_cls_fullname): lookup_type = AnyType(TypeOfAny.implementation_artifact) else: lookup_type = django_context.resolve_lookup_expected_type( ctx, model_cls, lookup_kwarg) # Managers as provided_type is not supported yet if isinstance(provided_type, Instance) and helpers.has_any_of_bases( provided_type.type, (fullnames.MANAGER_CLASS_FULLNAME, fullnames.QUERYSET_CLASS_FULLNAME)): return ctx.default_return_type helpers.check_types_compatible( ctx, expected_type=lookup_type, actual_type=provided_type, error_message=f"Incompatible type for lookup {lookup_kwarg!r}:", ) return ctx.default_return_type
def return_proper_field_type_from_get_field( ctx: MethodContext, django_context: DjangoContext) -> MypyType: # Options instance assert isinstance(ctx.type, Instance) # bail if list of generic params is empty if len(ctx.type.args) == 0: return ctx.default_return_type model_type = ctx.type.args[0] if not isinstance(model_type, Instance): return ctx.default_return_type model_cls = django_context.get_model_class_by_fullname( model_type.type.fullname) if model_cls is None: return ctx.default_return_type field_name_expr = helpers.get_call_argument_by_name(ctx, 'field_name') if field_name_expr is None: return ctx.default_return_type field_name = helpers.resolve_string_attribute_value( field_name_expr, django_context) if field_name is None: return ctx.default_return_type try: field = model_cls._meta.get_field(field_name) except FieldDoesNotExist as exc: # if model is abstract, do not raise exception, skip false positives if not model_cls._meta.abstract: ctx.api.fail(exc.args[0], ctx.context) return AnyType(TypeOfAny.from_error) field_fullname = helpers.get_class_fullname(field.__class__) return _get_field_instance(ctx, field_fullname)