Files
django-stubs/mypy_django_plugin/helpers.py

70 lines
2.2 KiB
Python

import typing
from typing import Dict, Optional, NamedTuple
from mypy.nodes import SymbolTableNode, Var, Expression
from mypy.plugin import FunctionContext
from mypy.types import Type, Instance, UnionType, NoneTyp
MODEL_CLASS_FULLNAME = 'django.db.models.base.Model'
QUERYSET_CLASS_FULLNAME = 'django.db.models.query.QuerySet'
FOREIGN_KEY_FULLNAME = 'django.db.models.fields.related.ForeignKey'
ONETOONE_FIELD_FULLNAME = 'django.db.models.fields.related.OneToOneField'
DUMMY_SETTINGS_BASE_CLASS = 'django.conf._DjangoConfLazyObject'
def create_new_symtable_node(name: str, kind: int, instance: Instance) -> SymbolTableNode:
new_var = Var(name, instance)
new_var.info = instance.type
return SymbolTableNode(kind, new_var,
plugin_generated=True)
Argument = NamedTuple('Argument', fields=[
('arg', Expression),
('arg_type', Type)
])
def get_call_signature_or_none(ctx: FunctionContext) -> Optional[Dict[str, Argument]]:
result: Dict[str, Argument] = {}
positional_args_only = []
positional_arg_types_only = []
for arg, arg_name, arg_type in zip(ctx.args, ctx.arg_names, ctx.arg_types):
if arg_name is None:
positional_args_only.append(arg)
positional_arg_types_only.append(arg_type)
continue
if len(arg) == 0 or len(arg_type) == 0:
continue
result[arg_name] = (arg[0], arg_type[0])
callee = ctx.context.callee
if '__init__' not in callee.node.names:
return None
init_type = callee.node.names['__init__'].type
arg_names = init_type.arg_names[1:]
for arg, arg_name, arg_type in zip(positional_args_only,
arg_names[:len(positional_args_only)],
positional_arg_types_only):
result[arg_name] = (arg[0], arg_type[0])
return result
def make_optional(typ: Type) -> Type:
return UnionType.make_simplified_union([typ, NoneTyp()])
def make_required(typ: Type) -> Type:
if not isinstance(typ, UnionType):
return typ
items = [item for item in typ.items if not isinstance(item, NoneTyp)]
return UnionType.make_union(items)
def get_obj_type_name(typ: typing.Type) -> str:
return typ.__module__ + '.' + typ.__qualname__