VimPlug/django-stubs
1
0
Fork 0
mirror of https://github.com/davidhalter/django-stubs.git synced 2025-12-07 04:34:29 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Specific return types for values and values list (#53)

Browse Source 
* 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
2019-03-25 10:53:09 +01:00
committed by Maxim Kurnikov
parent 5c6be7ad12
commit 5b455b729a
11 changed files with 648 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

97
mypy_django_plugin/helpers.py
View File

@@ -1,14 +1,16 @@
import typing
from typing import Dict, Optional
from collections import OrderedDict
from typing import Dict, Optional, cast
from mypy.checker import TypeChecker
from mypy.checker import TypeChecker, gen_unique_name
from mypy.mro import calculate_mro
from mypy.nodes import (
AssignmentStmt, ClassDef, Expression, ImportedName, Lvalue, MypyFile, NameExpr, SymbolNode, TypeInfo,
)
SymbolTable, SymbolTableNode, Block, GDEF, MDEF, Var)
from mypy.plugin import FunctionContext, MethodContext
from mypy.types import (
AnyType, Instance, NoneTyp, Type, TypeOfAny, TypeVarType, UnionType,
)
TupleType, TypedDictType)
MODEL_CLASS_FULLNAME = 'django.db.models.base.Model'
FIELD_FULLNAME = 'django.db.models.fields.Field'
@@ -211,7 +213,7 @@ def extract_field_setter_type(tp: Instance) -> Optional[Type]:
return None
def extract_field_getter_type(tp: Instance) -> Optional[Type]:
def extract_field_getter_type(tp: Type) -> Optional[Type]:
if not isinstance(tp, Instance):
return None
if tp.type.has_base(FIELD_FULLNAME):
@@ -235,6 +237,10 @@ def get_fields_metadata(model: TypeInfo) -> Dict[str, typing.Any]:
return get_django_metadata(model).setdefault('fields', {})
def get_lookups_metadata(model: TypeInfo) -> Dict[str, typing.Any]:
return get_django_metadata(model).setdefault('lookups', {})
def extract_explicit_set_type_of_model_primary_key(model: TypeInfo) -> Optional[Type]:
"""
If field with primary_key=True is set on the model, extract its __set__ type.
@@ -296,3 +302,84 @@ def get_assigned_value_for_class(type_info: TypeInfo, name: str) -> Optional[Exp
if isinstance(lvalue, NameExpr) and lvalue.name == name:
return rvalue
return None
def is_field_nullable(model: TypeInfo, field_name: str) -> bool:
return get_fields_metadata(model).get(field_name, {}).get('null', False)
def is_foreign_key(t: Type) -> bool:
if not isinstance(t, Instance):
return False
return has_any_of_bases(t.type, (FOREIGN_KEY_FULLNAME, ONETOONE_FIELD_FULLNAME))
def build_class_with_annotated_fields(api: TypeChecker, base: Type, fields: 'OrderedDict[str, Type]',
name: str) -> Instance:
"""Build an Instance with `name` that contains the specified `fields` as attributes and extends `base`."""
# Credit: This code is largely copied/modified from TypeChecker.intersect_instance_callable and
# NamedTupleAnalyzer.build_namedtuple_typeinfo
cur_module = cast(MypyFile, api.scope.stack[0])
gen_name = gen_unique_name(name, cur_module.names)
cdef = ClassDef(name, Block([]))
cdef.fullname = cur_module.fullname() + '.' + gen_name
info = TypeInfo(SymbolTable(), cdef, cur_module.fullname())
cdef.info = info
info.bases = [base]
def add_field(var: Var, is_initialized_in_class: bool = False,
is_property: bool = False) -> None:
var.info = info
var.is_initialized_in_class = is_initialized_in_class
var.is_property = is_property
var._fullname = '%s.%s' % (info.fullname(), var.name())
info.names[var.name()] = SymbolTableNode(MDEF, var)
vars = [Var(item, typ) for item, typ in fields.items()]
for var in vars:
add_field(var, is_property=True)
calculate_mro(info)
info.calculate_metaclass_type()
cur_module.names[gen_name] = SymbolTableNode(GDEF, info, plugin_generated=True)
return Instance(info, [])
def make_named_tuple(api: TypeChecker, fields: 'OrderedDict[str, Type]', name: str) -> Type:
if not fields:
# No fields specified, so fallback to a subclass of NamedTuple that allows
# __getattr__ / __setattr__ for any attribute name.
fallback = api.named_generic_type('django._NamedTupleAnyAttr', [])
else:
fallback = build_class_with_annotated_fields(
api=api,
base=api.named_generic_type('typing.NamedTuple', []),
fields=fields,
name=name
)
return TupleType(list(fields.values()), fallback=fallback)
def make_typeddict(api: TypeChecker, fields: 'OrderedDict[str, Type]', required_keys: typing.Set[str]) -> Type:
object_type = api.named_generic_type('mypy_extensions._TypedDict', [])
typed_dict_type = TypedDictType(fields, required_keys=required_keys, fallback=object_type)
return typed_dict_type
def make_tuple(api: TypeChecker, fields: typing.List[Type]) -> Type:
implicit_any = AnyType(TypeOfAny.special_form)
fallback = api.named_generic_type('builtins.tuple', [implicit_any])
return TupleType(fields, fallback=fallback)
def get_private_descriptor_type(type_info: TypeInfo, private_field_name: str, is_nullable: bool) -> Type:
node = type_info.get(private_field_name).node
if isinstance(node, Var):
descriptor_type = node.type
if is_nullable:
descriptor_type = make_optional(descriptor_type)
return descriptor_type
return AnyType(TypeOfAny.unannotated)

150
mypy_django_plugin/lookups.py Normal file
View File

@@ -0,0 +1,150 @@
import dataclasses
from typing import Union, List
from mypy.nodes import TypeInfo
from mypy.plugin import CheckerPluginInterface
from mypy.types import Type, Instance
from mypy_django_plugin import helpers
@dataclasses.dataclass
class RelatedModelNode:
typ: Instance
is_nullable: bool
@dataclasses.dataclass
class FieldNode:
typ: Type
LookupNode = Union[RelatedModelNode, FieldNode]
class LookupException(Exception):
pass
def resolve_lookup(api: CheckerPluginInterface, model_type_info: TypeInfo, lookup: str) -> List[LookupNode]:
"""Resolve a lookup str to a list of LookupNodes.
Each node represents a part of the lookup (separated by "__"), in order.
Each node is the Model or Field that was resolved.
Raises LookupException if there were any issues resolving the lookup.
"""
lookup_parts = lookup.split("__")
nodes = []
while lookup_parts:
lookup_part = lookup_parts.pop(0)
if not nodes:
current_node = None
else:
current_node = nodes[-1]
if current_node is None:
new_node = resolve_model_lookup(api, model_type_info, lookup_part)
elif isinstance(current_node, RelatedModelNode):
new_node = resolve_model_lookup(api, current_node.typ.type, lookup_part)
elif isinstance(current_node, FieldNode):
raise LookupException(f"Field lookups not yet supported for lookup {lookup}")
else:
raise LookupException(f"Unsupported node type: {type(current_node)}")
nodes.append(new_node)
return nodes
def resolve_model_lookup(api: CheckerPluginInterface, model_type_info: TypeInfo,
lookup: str) -> LookupNode:
"""Resolve a lookup on the given model."""
if lookup == 'pk':
# Primary keys are special-cased
primary_key_type = helpers.extract_primary_key_type_for_get(model_type_info)
if primary_key_type:
return FieldNode(primary_key_type)
else:
# No PK, use the get type for AutoField as PK type.
autofield_info = api.lookup_typeinfo('django.db.models.fields.AutoField')
pk_type = helpers.get_private_descriptor_type(autofield_info, '_pyi_private_get_type',
is_nullable=False)
return FieldNode(pk_type)
field_name = get_actual_field_name_for_lookup_field(lookup, model_type_info)
field_node = model_type_info.get(field_name)
if not field_node:
raise LookupException(
f'When resolving lookup "{lookup}", field "{field_name}" was not found in model {model_type_info.name()}')
if field_name.endswith('_id'):
field_name_without_id = field_name.rstrip('_id')
foreign_key_field = model_type_info.get(field_name_without_id)
if foreign_key_field is not None and helpers.is_foreign_key(foreign_key_field.type):
# Hack: If field ends with '_id' and there is a model field without the '_id' suffix, then use that field.
field_node = foreign_key_field
field_name = field_name_without_id
field_node_type = field_node.type
if field_node_type is None or not isinstance(field_node_type, Instance):
raise LookupException(
f'When resolving lookup "{lookup}", could not determine type for {model_type_info.name()}.{field_name}')
if helpers.is_foreign_key(field_node_type):
field_type = helpers.extract_field_getter_type(field_node_type)
is_nullable = helpers.is_optional(field_type)
if is_nullable:
field_type = helpers.make_required(field_type)
if isinstance(field_type, Instance):
return RelatedModelNode(typ=field_type, is_nullable=is_nullable)
else:
raise LookupException(f"Not an instance for field {field_type} lookup {lookup}")
field_type = helpers.extract_field_getter_type(field_node_type)
if field_type:
return FieldNode(typ=field_type)
else:
# Not a Field
if field_name == 'id':
# If no 'id' field was fouond, use an int
return FieldNode(api.named_generic_type('builtins.int', []))
related_manager_arg = None
if field_node_type.type.has_base(helpers.RELATED_MANAGER_CLASS_FULLNAME):
related_manager_arg = field_node_type.args[0]
if related_manager_arg is not None:
# Reverse relation
return RelatedModelNode(typ=related_manager_arg, is_nullable=True)
raise LookupException(
f'When resolving lookup "{lookup}", could not determine type for {model_type_info.name()}.{field_name}')
def get_actual_field_name_for_lookup_field(lookup: str, model_type_info: TypeInfo) -> str:
"""Attempt to find out the real field name if this lookup is a related_query_name (for reverse relations).
If it's not, return the original lookup.
"""
lookups_metadata = helpers.get_lookups_metadata(model_type_info)
lookup_metadata = lookups_metadata.get(lookup)
if lookup_metadata is None:
# If not found on current model, look in all bases for their lookup metadata
for base in model_type_info.mro:
lookups_metadata = helpers.get_lookups_metadata(base)
lookup_metadata = lookups_metadata.get(lookup)
if lookup_metadata:
break
if not lookup_metadata:
lookup_metadata = {}
related_name = lookup_metadata.get('related_query_name_target', None)
if related_name:
# If the lookup is a related lookup, then look at the field specified by related_name.
# This is to support if related_query_name is set and differs from.
field_name = related_name
else:
field_name = lookup
return field_name

15
mypy_django_plugin/main.py
View File

@@ -1,6 +1,5 @@
from functools import partial
import os
from functools import partial
from typing import Callable, Dict, Optional, Union, cast
from mypy.checker import TypeChecker
@@ -23,6 +22,7 @@ from mypy_django_plugin.transformers.migrations import (
determine_model_cls_from_string_for_migrations, get_string_value_from_expr,
)
from mypy_django_plugin.transformers.models import process_model_class
from mypy_django_plugin.transformers.queryset import extract_proper_type_for_values_and_values_list
from mypy_django_plugin.transformers.settings import (
AddSettingValuesToDjangoConfObject, get_settings_metadata,
)
@@ -165,7 +165,7 @@ def extract_and_return_primary_key_of_bound_related_field_parameter(ctx: Attribu
if primary_key_type:
return primary_key_type
is_nullable = helpers.get_fields_metadata(ctx.type.type).get(field_name, {}).get('null', False)
is_nullable = helpers.is_field_nullable(ctx.type.type, field_name)
if is_nullable:
return helpers.make_optional(ctx.default_attr_type)
@@ -292,7 +292,10 @@ class DjangoPlugin(Plugin):
if self.django_settings_module:
settings_modules.append(self.django_settings_module)
monkeypatch.add_modules_as_a_source_seed_files(settings_modules)
auto_imports = ['mypy_extensions']
auto_imports.extend(settings_modules)
monkeypatch.add_modules_as_a_source_seed_files(auto_imports)
monkeypatch.inject_modules_as_dependencies_for_django_conf_settings(settings_modules)
def _get_current_model_bases(self) -> Dict[str, int]:
@@ -359,10 +362,10 @@ class DjangoPlugin(Plugin):
if sym and isinstance(sym.node, TypeInfo) and sym.node.has_base(helpers.FORM_MIXIN_CLASS_FULLNAME):
return extract_proper_type_for_get_form
if method_name == 'values_list':
if method_name in ('values', 'values_list'):
sym = self.lookup_fully_qualified(class_name)
if sym and isinstance(sym.node, TypeInfo) and sym.node.has_base(helpers.QUERYSET_CLASS_FULLNAME):
return extract_proper_type_for_values_list
return partial(extract_proper_type_for_values_and_values_list, method_name)
if fullname in {'django.apps.registry.Apps.get_model',
'django.db.migrations.state.StateApps.get_model'}:

22
mypy_django_plugin/transformers/fields.py
View File

@@ -1,10 +1,10 @@
from typing import Optional, cast
from mypy.checker import TypeChecker
from mypy.nodes import ListExpr, NameExpr, StrExpr, TupleExpr, TypeInfo, Var
from mypy.nodes import ListExpr, NameExpr, StrExpr, TupleExpr, TypeInfo
from mypy.plugin import FunctionContext
from mypy.types import (
AnyType, CallableType, Instance, TupleType, Type, TypeOfAny, UnionType,
AnyType, CallableType, Instance, TupleType, Type, UnionType,
)
from mypy_django_plugin import helpers
@@ -88,23 +88,13 @@ def fill_descriptor_types_for_related_field(ctx: FunctionContext) -> Type:
return helpers.reparametrize_instance(ctx.default_return_type, new_args=args)
def get_private_descriptor_type(type_info: TypeInfo, private_field_name: str, is_nullable: bool) -> Type:
node = type_info.get(private_field_name).node
if isinstance(node, Var):
descriptor_type = node.type
if is_nullable:
descriptor_type = helpers.make_optional(descriptor_type)
return descriptor_type
return AnyType(TypeOfAny.unannotated)
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 = 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)
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])

9
mypy_django_plugin/transformers/init_create.py
View File

@@ -6,7 +6,6 @@ from mypy.plugin import FunctionContext, MethodContext
from mypy.types import AnyType, Instance, Type, TypeOfAny
from mypy_django_plugin import helpers
from mypy_django_plugin.transformers.fields import get_private_descriptor_type
def extract_base_pointer_args(model: TypeInfo) -> Set[str]:
@@ -162,8 +161,8 @@ def extract_expected_types(ctx: FunctionContext, model: TypeInfo,
if not pk_type:
# extract set type of AutoField
autofield_info = api.lookup_typeinfo('django.db.models.fields.AutoField')
pk_type = get_private_descriptor_type(autofield_info, '_pyi_private_set_type',
is_nullable=is_nullable)
pk_type = helpers.get_private_descriptor_type(autofield_info, '_pyi_private_set_type',
is_nullable=is_nullable)
related_primary_key_type = pk_type
if is_init:
@@ -185,8 +184,8 @@ def extract_expected_types(ctx: FunctionContext, model: TypeInfo,
# if CharField(blank=True,...) and not nullable, then field can be None in __init__
elif (
helpers.has_any_of_bases(typ.type, (helpers.CHAR_FIELD_FULLNAME,)) and is_init and
field_metadata.get('blank', False) and not field_metadata.get('null', False)
helpers.has_any_of_bases(typ.type, (helpers.CHAR_FIELD_FULLNAME,)) and is_init and
field_metadata.get('blank', False) and not field_metadata.get('null', False)
):
field_type = helpers.make_optional(field_type)

28
mypy_django_plugin/transformers/models.py
View File

@@ -172,17 +172,35 @@ class AddRelatedManagers(ModelClassInitializer):
ref_to_fullname = module_name + '.' + exc.model_cls_name
if self.model_classdef.fullname == ref_to_fullname:
related_manager_name = defn.name.lower() + '_set'
related_name = defn.name.lower() + '_set'
if 'related_name' in rvalue.arg_names:
related_name_expr = rvalue.args[rvalue.arg_names.index('related_name')]
if not isinstance(related_name_expr, StrExpr):
return None
related_manager_name = related_name_expr.value
continue
related_name = related_name_expr.value
if related_name == '+':
# No backwards relation is desired
continue
if 'related_query_name' in rvalue.arg_names:
related_query_name_expr = rvalue.args[rvalue.arg_names.index('related_query_name')]
if not isinstance(related_query_name_expr, StrExpr):
related_query_name = None
else:
related_query_name = related_query_name_expr.value
# TODO: Handle defaulting to model name if related_name is not set
else:
# No related_query_name specified, default to related_name
related_query_name = related_name
typ = get_related_field_type(rvalue, self.api, defn.info)
if typ is None:
return None
self.add_new_node_to_model_class(related_manager_name, typ)
continue
self.add_new_node_to_model_class(related_name, typ)
if related_query_name is not None:
# Only create related_query_name if it is a string literal
helpers.get_lookups_metadata(self.model_classdef.info)[related_query_name] = {
'related_query_name_target': related_name
}
def iter_over_classdefs(module_file: MypyFile) -> Iterator[ClassDef]:

139
mypy_django_plugin/transformers/queryset.py Normal file
View 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