def extract_proper_type_for_values_list(ctx: MethodContext) -> Type: object_type = ctx.type if not isinstance(object_type, Instance): return ctx.default_return_type flat = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'flat')) named = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'named')) ret = ctx.default_return_type any_type = AnyType(TypeOfAny.implementation_artifact) if named and flat: ctx.api.fail("'flat' and 'named' can't be used together.", ctx.context) return ret elif named: # TODO: Fill in namedtuple fields/types row_arg = ctx.api.named_generic_type('typing.NamedTuple', []) elif flat: # TODO: Figure out row_arg type dependent on the argument passed in if len(ctx.args[0]) > 1: ctx.api.fail( "'flat' is not valid when values_list is called with more than one field.", ctx.context) return ret row_arg = any_type else: # TODO: Figure out tuple argument types dependent on the arguments passed in row_arg = ctx.api.named_generic_type('builtins.tuple', [any_type]) first_arg = ret.args[0] if len(ret.args) > 0 else any_type new_type_args = [first_arg, row_arg] return helpers.reparametrize_instance(ret, new_type_args)
def record_field_properties_into_outer_model_class( ctx: FunctionContext) -> None: api = cast(TypeChecker, ctx.api) outer_model = api.scope.active_class() if outer_model is None or not outer_model.has_base( helpers.MODEL_CLASS_FULLNAME): # outside models.Model class, undetermined return field_name = None for name_expr, stmt in helpers.iter_over_assignments(outer_model.defn): if stmt == ctx.context and isinstance(name_expr, NameExpr): field_name = name_expr.name break if field_name is None: return fields_metadata = outer_model.metadata.setdefault('django', {}).setdefault( 'fields', {}) # primary key is_primary_key = False primary_key_arg = helpers.get_argument_by_name(ctx, 'primary_key') if primary_key_arg: is_primary_key = helpers.parse_bool(primary_key_arg) fields_metadata[field_name] = {'primary_key': is_primary_key} # choices choices_arg = helpers.get_argument_by_name(ctx, 'choices') if choices_arg and isinstance(choices_arg, (TupleExpr, ListExpr)): # iterable of 2 element tuples of two kinds _, analyzed_choices = api.analyze_iterable_item_type(choices_arg) if isinstance(analyzed_choices, TupleType): first_element_type = analyzed_choices.items[0] if isinstance(first_element_type, Instance): fields_metadata[field_name][ 'choices'] = first_element_type.type.fullname() # nullability null_arg = helpers.get_argument_by_name(ctx, 'null') is_nullable = False if null_arg: is_nullable = helpers.parse_bool(null_arg) fields_metadata[field_name]['null'] = is_nullable # is_blankable blank_arg = helpers.get_argument_by_name(ctx, 'blank') is_blankable = False if blank_arg: is_blankable = helpers.parse_bool(blank_arg) fields_metadata[field_name]['blank'] = is_blankable # default default_arg = helpers.get_argument_by_name(ctx, 'default') if default_arg and not helpers.is_none_expr(default_arg): fields_metadata[field_name]['default_specified'] = True
def set_descriptor_types_for_field(ctx: FunctionContext) -> Instance: default_return_type = cast(Instance, ctx.default_return_type) is_nullable = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'null')) if not is_nullable and default_return_type.type.has_base(helpers.CHAR_FIELD_FULLNAME): # blank=True for CharField can be interpreted as null=True is_nullable = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'blank')) set_type = get_private_descriptor_type(default_return_type.type, '_pyi_private_set_type', is_nullable=is_nullable) get_type = get_private_descriptor_type(default_return_type.type, '_pyi_private_get_type', is_nullable=is_nullable) return helpers.reparametrize_instance(default_return_type, [set_type, get_type])
def set_descriptor_types_for_field(ctx: FunctionContext) -> Instance: default_return_type = cast(Instance, ctx.default_return_type) is_nullable = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'null')) set_type = helpers.get_private_descriptor_type(default_return_type.type, '_pyi_private_set_type', is_nullable=is_nullable) get_type = helpers.get_private_descriptor_type(default_return_type.type, '_pyi_private_get_type', is_nullable=is_nullable) return helpers.reparametrize_instance(default_return_type, [set_type, get_type])
def extract_proper_type_queryset_values_list(ctx: MethodContext) -> Type: object_type = ctx.type if not isinstance(object_type, Instance): return ctx.default_return_type ret = ctx.default_return_type model_arg = get_queryset_model_arg(ctx.default_return_type) # model_arg: Union[AnyType, Type] = ret.args[0] if len(ret.args) > 0 else any_type column_names: List[Optional[str]] = [] column_types: OrderedDict[str, Type] = OrderedDict() fields_arg_expr = ctx.args[ctx.callee_arg_names.index('fields')] fields_param_is_specified = True if len(fields_arg_expr) == 0: # values_list/values with no args is not yet supported, so default to Any types for field types # It should in the future include all model fields, "extra" fields and "annotated" fields fields_param_is_specified = False if isinstance(model_arg, Instance): model_type_info = model_arg.type else: model_type_info = None any_type = AnyType(TypeOfAny.implementation_artifact) # Figure out each field name passed to fields only_strings_as_fields_expressions = True for field_expr in fields_arg_expr: if isinstance(field_expr, StrExpr): field_name = field_expr.value column_names.append(field_name) # Default to any type column_types[field_name] = any_type if model_type_info: resolved_lookup_type = resolve_values_lookup( ctx.api, model_type_info, field_name) if resolved_lookup_type is not None: column_types[field_name] = resolved_lookup_type else: # Dynamic field names are partially supported for values_list, but not values column_names.append(None) only_strings_as_fields_expressions = False flat = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'flat')) named = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'named')) api = cast(TypeChecker, ctx.api) if named and flat: api.fail("'flat' and 'named' can't be used together.", ctx.context) return ret elif named: # named=True, flat=False -> List[NamedTuple] if fields_param_is_specified and only_strings_as_fields_expressions: row_arg = helpers.make_named_tuple(api, fields=column_types, name="Row") else: # fallback to catch-all NamedTuple row_arg = helpers.make_named_tuple(api, fields=OrderedDict(), name="Row") elif flat: # named=False, flat=True -> List of elements if len(ctx.args[0]) > 1: api.fail( "'flat' is not valid when values_list is called with more than one field.", ctx.context) return ctx.default_return_type if fields_param_is_specified and only_strings_as_fields_expressions: # Grab first element row_arg = column_types[column_names[0]] else: row_arg = any_type else: # named=False, flat=False -> List[Tuple] if fields_param_is_specified: args = [ # Fallback to Any if the column name is unknown (e.g. dynamic) column_types.get(column_name, any_type) if column_name is not None else any_type for column_name in column_names ] else: args = [any_type] row_arg = helpers.make_tuple(api, fields=args) new_type_args = [model_arg, row_arg] return helpers.reparametrize_instance(ret, new_type_args)
def extract_proper_type_for_values_and_values_list(method_name: str, ctx: MethodContext) -> Type: api = cast(TypeChecker, ctx.api) object_type = ctx.type if not isinstance(object_type, Instance): return ctx.default_return_type ret = ctx.default_return_type any_type = AnyType(TypeOfAny.implementation_artifact) fields_arg_expr = ctx.args[ctx.callee_arg_names.index('fields')] model_arg: Union[AnyType, Type] = ret.args[0] if len(ret.args) > 0 else any_type column_names: List[Optional[str]] = [] column_types: OrderedDict[str, Type] = OrderedDict() fill_column_types = True if len(fields_arg_expr) == 0: # values_list/values with no args is not yet supported, so default to Any types for field types # It should in the future include all model fields, "extra" fields and "annotated" fields fill_column_types = False if isinstance(model_arg, Instance): model_type_info = model_arg.type else: model_type_info = None # Figure out each field name passed to fields has_dynamic_column_names = False for field_expr in fields_arg_expr: if isinstance(field_expr, StrExpr): field_name = field_expr.value column_names.append(field_name) # Default to any type column_types[field_name] = any_type if model_type_info: resolved_lookup_type = resolve_values_lookup( ctx.api, model_type_info, field_name) if resolved_lookup_type is not None: column_types[field_name] = resolved_lookup_type else: # Dynamic field names are partially supported for values_list, but not values column_names.append(None) has_dynamic_column_names = True if method_name == 'values_list': flat = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'flat')) named = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'named')) if named and flat: api.fail("'flat' and 'named' can't be used together.", ctx.context) return ret elif named: if fill_column_types and not has_dynamic_column_names: row_arg = helpers.make_named_tuple(api, fields=column_types, name="Row") else: row_arg = helpers.make_named_tuple(api, fields=OrderedDict(), name="Row") elif flat: if len(ctx.args[0]) > 1: api.fail( "'flat' is not valid when values_list is called with more than one field.", ctx.context) return ret if fill_column_types and not has_dynamic_column_names: # Grab first element row_arg = column_types[column_names[0]] else: row_arg = any_type else: if fill_column_types: args = [ # Fallback to Any if the column name is unknown (e.g. dynamic) column_types.get(column_name, any_type) if column_name is not None else any_type for column_name in column_names ] else: args = [any_type] row_arg = helpers.make_tuple(api, fields=args) elif method_name == 'values': expression_arg_names = ctx.arg_names[ctx.callee_arg_names.index( 'expressions')] for expression_name in expression_arg_names: # Arbitrary additional annotation expressions are supported, but they all have type Any for now column_names.append(expression_name) column_types[expression_name] = any_type if fill_column_types and not has_dynamic_column_names: row_arg = helpers.make_typeddict(api, fields=column_types, required_keys=set()) else: return ctx.default_return_type else: raise Exception( f"extract_proper_type_for_values_list doesn't support method {method_name}" ) new_type_args = [model_arg, row_arg] return helpers.reparametrize_instance(ret, new_type_args)