mirror of
https://github.com/davidhalter/django-stubs.git
synced 2025-12-08 21:14:49 +08:00
Specific return types for values and values list (#53)
* Instead of using Literal types, overload QuerySet.values_list in the plugin. Fixes #43. - Add a couple of extra type checks that Django makes: 1) 'flat' and 'named' can't be used together. 2) 'flat' is not valid when values_list is called with more than one field. * Determine better row types for values_list/values based on fields specified. - In the case of values_list, we use a Row type with either a single primitive, Tuple, or NamedTuple. - In the case of values, we use a TypedDict. - In both cases, Any is used as a fallback for individual fields if those fields cannot be resolved. A couple other fixes I made along the way: - Don't create reverse relation for ForeignKeys with related_name='+' - Don't skip creating other related managers in AddRelatedManagers if a dynamic value is encountered for related_name parameter, or if the type cannot be determined. * Fix for TypedDict so that they are considered anonymous. * Clean up some comments. * Implement making TypedDict anonymous in a way that doesn't crash sometimes. * Fix flake8 errors. * Remove even uglier hack about making TypedDict anonymous. * Address review comments. Write a few better comments inside tests. * Fix crash when running with mypyc ("interpreted classes cannot inherit from compiled") due to the way I extended TypedDictType. - Implemented the hack in another way that works on mypyc. - Added a couple extra tests of accessing 'id' / 'pk' via values_list. * Fix flake8 errors. * Support annotation expressions (use type Any) for TypedDicts row types returned by values_list. - Bonus points: handle values_list gracefully (use type Any) where Tuples are returned where some of the fields arguments are not string literals.
This commit is contained in:
Seth Yastrov
committed by
Maxim Kurnikov
Maxim Kurnikov
parent
5c6be7ad12
commit
5b455b729a
139
mypy_django_plugin/transformers/queryset.py
Normal file
139
mypy_django_plugin/transformers/queryset.py
Normal file
@@ -0,0 +1,139 @@
|
||||
from collections import OrderedDict
|
||||
from typing import Union, List, cast, Optional
|
||||
|
||||
from mypy.checker import TypeChecker
|
||||
from mypy.nodes import StrExpr, TypeInfo
|
||||
from mypy.plugin import MethodContext, CheckerPluginInterface
|
||||
from mypy.types import Type, Instance, AnyType, TypeOfAny
|
||||
|
||||
from mypy_django_plugin import helpers
|
||||
from mypy_django_plugin.lookups import resolve_lookup, RelatedModelNode, LookupException
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
def resolve_values_lookup(api: CheckerPluginInterface, model_type_info: TypeInfo, lookup: str) -> Optional[Type]:
|
||||
"""Resolves a values/values_list lookup if possible, to a Type."""
|
||||
try:
|
||||
nodes = resolve_lookup(api, model_type_info, lookup)
|
||||
except LookupException:
|
||||
nodes = []
|
||||
|
||||
if not nodes:
|
||||
return None
|
||||
|
||||
make_optional = False
|
||||
|
||||
for node in nodes:
|
||||
if isinstance(node, RelatedModelNode) and node.is_nullable:
|
||||
# All lookups following a relation which is nullable should be optional
|
||||
make_optional = True
|
||||
|
||||
node = nodes[-1]
|
||||
|
||||
node_type = node.typ
|
||||
if isinstance(node, RelatedModelNode):
|
||||
# Related models used in values/values_list get resolved to the primary key of the related model.
|
||||
# So, we lookup the pk of that model.
|
||||
pk_lookup_nodes = resolve_lookup(api, node_type.type, "pk")
|
||||
if not pk_lookup_nodes:
|
||||
return None
|
||||
node_type = pk_lookup_nodes[0].typ
|
||||
if make_optional:
|
||||
return helpers.make_optional(node_type)
|
||||
else:
|
||||
return node_type
|
||||
Reference in New Issue
Block a user