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

fix couple edge cases with __init__

Browse Source
This commit is contained in:
Maxim Kurnikov
2019-02-10 04:32:27 +03:00
parent 5f6f597266
commit 6b7507206a
8 changed files with 172 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
mypy_django_plugin/helpers.py
View File

@@ -1,9 +1,10 @@
import typing
from typing import Dict, Optional
from mypy.nodes import Expression, FuncDef, ImportedName, MypyFile, NameExpr, SymbolNode, TypeInfo, Var
from mypy.nodes import Expression, FuncDef, ImportedName, MypyFile, NameExpr, SymbolNode, TypeInfo, Var, AssignmentStmt, \
CallExpr
from mypy.plugin import FunctionContext
from mypy.types import AnyType, CallableType, Instance, Type, TypeOfAny, TypeVarType
from mypy.types import AnyType, CallableType, Instance, Type, TypeOfAny, TypeVarType, UnionType
MODEL_CLASS_FULLNAME = 'django.db.models.base.Model'
FIELD_FULLNAME = 'django.db.models.fields.Field'
@@ -101,10 +102,23 @@ def fill_typevars_with_any(instance: Instance) -> Type:
return reparametrize_with(instance, [AnyType(TypeOfAny.unannotated)])
def extract_typevar_value(tp: Instance, typevar_name: str):
def extract_typevar_value(tp: Instance, typevar_name: str) -> Type:
if typevar_name in {'_T', '_T_co'}:
if '_T' in tp.type.type_vars:
return tp.args[tp.type.type_vars.index('_T')]
if '_T_co' in tp.type.type_vars:
return tp.args[tp.type.type_vars.index('_T_co')]
return tp.args[tp.type.type_vars.index(typevar_name)]
def fill_typevars(tp: Instance, type_to_fill: Instance) -> Instance:
typevar_values: typing.List[Type] = []
for typevar_arg in type_to_fill.args:
if isinstance(typevar_arg, TypeVarType):
typevar_values.append(extract_typevar_value(tp, typevar_arg.name))
return reparametrize_with(type_to_fill, typevar_values)
def extract_field_setter_type(tp: Instance) -> Optional[Type]:
if tp.type.has_base(FIELD_FULLNAME):
set_method = tp.type.get_method('__set__')
@@ -112,13 +126,15 @@ def extract_field_setter_type(tp: Instance) -> Optional[Type]:
if 'value' in set_method.type.arg_names:
set_value_type = set_method.type.arg_types[set_method.type.arg_names.index('value')]
if isinstance(set_value_type, Instance):
typevar_values: typing.List[Type] = []
for typevar_arg in set_value_type.args:
if isinstance(typevar_arg, TypeVarType):
typevar_values.append(extract_typevar_value(tp, typevar_arg.name))
# if there are typevars, extract from
set_value_type = reparametrize_with(set_value_type, typevar_values)
set_value_type = fill_typevars(tp, set_value_type)
return set_value_type
elif isinstance(set_value_type, UnionType):
items_no_typevars = []
for item in set_value_type.items:
if isinstance(item, Instance):
item = fill_typevars(tp, item)
items_no_typevars.append(item)
return UnionType(items_no_typevars)
get_method = tp.type.get_method('__get__')
if isinstance(get_method, FuncDef) and isinstance(get_method.type, CallableType):
@@ -131,31 +147,20 @@ def extract_field_setter_type(tp: Instance) -> Optional[Type]:
def extract_primary_key_type(model: TypeInfo) -> Optional[Type]:
# only primary keys defined in current class for now
for sym in model.names.values():
if isinstance(sym.node, Var) and isinstance(sym.node.type, Instance):
tp = sym.node.type
if tp.type.metadata.get('django', {}).get('defined_as_primary_key'):
field_type = extract_field_setter_type(tp)
return field_type
for stmt in model.defn.defs.body:
if isinstance(stmt, AssignmentStmt) and isinstance(stmt.rvalue, CallExpr):
name_expr = stmt.lvalues[0]
if isinstance(name_expr, NameExpr):
name = name_expr.name
if 'primary_key' in stmt.rvalue.arg_names:
is_primary_key = stmt.rvalue.args[stmt.rvalue.arg_names.index('primary_key')]
if is_primary_key:
return extract_field_setter_type(model.names[name].type)
return None
def extract_expected_types(ctx: FunctionContext, model: TypeInfo) -> Dict[str, Type]:
expected_types: Dict[str, Type] = {}
for base in model.mro:
for name, sym in base.names.items():
if isinstance(sym.node, Var) and isinstance(sym.node.type, Instance):
tp = sym.node.type
field_type = extract_field_setter_type(tp)
if tp.type.fullname() == FOREIGN_KEY_FULLNAME:
ref_to_model = tp.args[0]
if isinstance(ref_to_model, Instance) and ref_to_model.type.has_base(MODEL_CLASS_FULLNAME):
primary_key_type = extract_primary_key_type(ref_to_model.type)
if not primary_key_type:
primary_key_type = AnyType(TypeOfAny.special_form)
expected_types[name + '_id'] = primary_key_type
if field_type:
expected_types[name] = field_type
primary_key_type = extract_primary_key_type(model)
if not primary_key_type:
@@ -164,6 +169,21 @@ def extract_expected_types(ctx: FunctionContext, model: TypeInfo) -> Dict[str, T
expected_types['id'] = ctx.api.named_generic_type('builtins.int', [])
expected_types['pk'] = primary_key_type
for base in model.mro:
for name, sym in base.names.items():
if isinstance(sym.node, Var) and isinstance(sym.node.type, Instance):
tp = sym.node.type
field_type = extract_field_setter_type(tp)
if tp.type.fullname() in {FOREIGN_KEY_FULLNAME, ONETOONE_FIELD_FULLNAME}:
ref_to_model = tp.args[0]
if isinstance(ref_to_model, Instance) and ref_to_model.type.has_base(MODEL_CLASS_FULLNAME):
primary_key_type = extract_primary_key_type(ref_to_model.type)
if not primary_key_type:
primary_key_type = AnyType(TypeOfAny.special_form)
expected_types[name + '_id'] = primary_key_type
if field_type:
expected_types[name] = field_type
return expected_types