mirror of
https://github.com/davidhalter/django-stubs.git
synced 2025-12-06 20:24:31 +08:00
42 lines
1.5 KiB
Python
42 lines
1.5 KiB
Python
from typing import Iterator, List, cast
|
|
|
|
from mypy.nodes import ClassDef, AssignmentStmt, CallExpr
|
|
from mypy.plugin import FunctionContext, ClassDefContext
|
|
from mypy.semanal import SemanticAnalyzerPass2
|
|
from mypy.types import Type, Instance
|
|
|
|
from mypy_django_plugin.plugins.related_fields import add_new_var_node_to_class
|
|
|
|
|
|
def determine_type_of_array_field(ctx: FunctionContext) -> Type:
|
|
if 'base_field' not in ctx.arg_names:
|
|
return ctx.default_return_type
|
|
|
|
base_field_arg_type = ctx.arg_types[ctx.arg_names.index('base_field')][0]
|
|
return ctx.api.named_generic_type(ctx.context.callee.fullname,
|
|
args=[base_field_arg_type.type.names['__get__'].type.ret_type])
|
|
|
|
|
|
def get_assignments(klass: ClassDef) -> List[AssignmentStmt]:
|
|
stmts = []
|
|
for stmt in klass.defs.body:
|
|
if not isinstance(stmt, AssignmentStmt):
|
|
continue
|
|
if len(stmt.lvalues) > 1:
|
|
# not supported yet
|
|
continue
|
|
stmts.append(stmt)
|
|
return stmts
|
|
|
|
|
|
def add_int_id_attribute_if_primary_key_true_is_not_present(ctx: ClassDefContext) -> None:
|
|
api = cast(SemanticAnalyzerPass2, ctx.api)
|
|
for stmt in get_assignments(ctx.cls):
|
|
if (isinstance(stmt.rvalue, CallExpr)
|
|
and 'primary_key' in stmt.rvalue.arg_names
|
|
and api.parse_bool(stmt.rvalue.args[stmt.rvalue.arg_names.index('primary_key')])):
|
|
break
|
|
else:
|
|
add_new_var_node_to_class(ctx.cls.info, 'id', api.builtin_type('builtins.int'))
|
|
|