Skip to content

Commit 838942c

Browse files
mkurnikovmkurnikov
committed
wip
1 parent 540e28f commit 838942c

File tree

7 files changed

+181
-37
lines changed

7 files changed

+181
-37
lines changed

‎.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,3 +11,4 @@ pip-wheel-metadata/
1111
.pytest_cache/
1212
/.envrc
1313
/.direnv
14+
django-sources/

‎django-stubs/contrib/auth/base_user.pyi

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ class AbstractBaseUser(models.Model):
1616
password = models.CharField(max_length=128)
1717
last_login = models.DateTimeField(blank=True, null=True)
1818

19+
is_active: bool = ...
20+
1921
REQUIRED_FIELDS: List[str] = ...
2022
class Meta: ...
2123
def get_username(self) -> str: ...

‎mypy_django_plugin/lib/helpers.py

Lines changed: 72 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
11
from collections import OrderedDict
22
from typing import (
3-
TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Optional, Set, Union, cast,
4-
)
3+
TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Optional, Set, Union, Tuple)
54

6-
from django.db.models.fields import Field
75
from django.db.models.fields.related import RelatedField
86
from django.db.models.fields.reverse_related import ForeignObjectRel
97
from mypy import checker
@@ -12,14 +10,17 @@
1210
from mypy.nodes import (
1311
GDEF, MDEF, Block, ClassDef, Expression, MemberExpr, MypyFile, NameExpr, StrExpr, SymbolNode, SymbolTable,
1412
SymbolTableNode, TypeInfo, Var,
15-
)
13+
FuncDef, Argument)
1614
from mypy.plugin import (
1715
AttributeContext, CheckerPluginInterface, FunctionContext, MethodContext,
18-
)
19-
from mypy.types import AnyType, Instance, NoneTyp, TupleType
16+
ClassDefContext, DynamicClassDefContext)
17+
from mypy.plugins.common import add_method
18+
from mypy.semanal import SemanticAnalyzer
19+
from mypy.types import AnyType, Instance, NoneTyp, TupleType, CallableType
2020
from mypy.types import Type as MypyType
2121
from mypy.types import TypedDictType, TypeOfAny, UnionType
2222

23+
from django.db.models.fields import Field
2324
from mypy_django_plugin.lib import fullnames
2425

2526
if TYPE_CHECKING:
@@ -174,7 +175,7 @@ def get_nested_meta_node_for_current_class(info: TypeInfo) -> Optional[TypeInfo]
174175

175176

176177
def add_new_class_for_module(module: MypyFile, name: str, bases: List[Instance],
177-
fields: 'OrderedDict[str, MypyType]') -> TypeInfo:
178+
fields: 'Optional[OrderedDict[str, MypyType]]' = None) -> TypeInfo:
178179
new_class_unique_name = checker.gen_unique_name(name, module.names)
179180

180181
# make new class expression
@@ -188,17 +189,44 @@ def add_new_class_for_module(module: MypyFile, name: str, bases: List[Instance],
188189
new_typeinfo.calculate_metaclass_type()
189190

190191
# add fields
191-
for field_name, field_type in fields.items():
192-
var = Var(field_name, type=field_type)
193-
var.info = new_typeinfo
194-
var._fullname = new_typeinfo.fullname + '.' + field_name
195-
new_typeinfo.names[field_name] = SymbolTableNode(MDEF, var, plugin_generated=True)
192+
if fields:
193+
for field_name, field_type in fields.items():
194+
var = Var(field_name, type=field_type)
195+
var.info = new_typeinfo
196+
var._fullname = new_typeinfo.fullname + '.' + field_name
197+
new_typeinfo.names[field_name] = SymbolTableNode(MDEF, var, plugin_generated=True)
196198

197199
classdef.info = new_typeinfo
198200
module.names[new_class_unique_name] = SymbolTableNode(GDEF, new_typeinfo, plugin_generated=True)
199201
return new_typeinfo
200202

201203

204+
# def add_new_class_for_module_semanal(semanal: SemanticAnalyzer, module: MypyFile, name: str, bases: List[Instance],
205+
# fields: 'OrderedDict[str, MypyType]') -> TypeInfo:
206+
# new_class_unique_name = checker.gen_unique_name(name, module.names)
207+
#
208+
# # make new class expression
209+
# classdef = ClassDef(new_class_unique_name, Block([]))
210+
# classdef.fullname = module.fullname + '.' + new_class_unique_name
211+
#
212+
# # make new TypeInfo
213+
# new_typeinfo = TypeInfo(SymbolTable(), classdef, module.fullname)
214+
# new_typeinfo.bases = bases
215+
# calculate_mro(new_typeinfo)
216+
# new_typeinfo.calculate_metaclass_type()
217+
#
218+
# # add fields
219+
# for field_name, field_type in fields.items():
220+
# var = Var(field_name, type=field_type)
221+
# var.info = new_typeinfo
222+
# var._fullname = new_typeinfo.fullname + '.' + field_name
223+
# new_typeinfo.names[field_name] = SymbolTableNode(MDEF, var, plugin_generated=True)
224+
#
225+
# classdef.info = new_typeinfo
226+
# module.names[new_class_unique_name] = SymbolTableNode(GDEF, new_typeinfo, plugin_generated=True)
227+
# return new_typeinfo
228+
229+
202230
def get_current_module(api: TypeChecker) -> MypyFile:
203231
current_module = None
204232
for item in reversed(api.scope.stack):
@@ -269,10 +297,16 @@ def resolve_string_attribute_value(attr_expr: Expression, ctx: Union[FunctionCon
269297
return None
270298

271299

300+
def get_semanal_api(ctx: Union[ClassDefContext, DynamicClassDefContext]) -> SemanticAnalyzer:
301+
if not isinstance(ctx.api, SemanticAnalyzer):
302+
raise ValueError('Not a SemanticAnalyzer')
303+
return ctx.api
304+
305+
272306
def get_typechecker_api(ctx: Union[AttributeContext, MethodContext, FunctionContext]) -> TypeChecker:
273307
if not isinstance(ctx.api, TypeChecker):
274308
raise ValueError('Not a TypeChecker')
275-
return cast(TypeChecker, ctx.api)
309+
return ctx.api
276310

277311

278312
def is_model_subclass_info(info: TypeInfo, django_context: 'DjangoContext') -> bool:
@@ -298,3 +332,28 @@ def add_new_sym_for_info(info: TypeInfo, *, name: str, sym_type: MypyType) -> No
298332
var.is_inferred = True
299333
info.names[name] = SymbolTableNode(MDEF, var,
300334
plugin_generated=True)
335+
336+
337+
def _prepare_new_method_arguments(node: FuncDef) -> Tuple[List[Argument], MypyType]:
338+
arguments = []
339+
for argument in node.arguments[1:]:
340+
if argument.type_annotation is None:
341+
argument.type_annotation = AnyType(TypeOfAny.unannotated)
342+
arguments.append(argument)
343+
344+
if isinstance(node.type, CallableType):
345+
return_type = node.type.ret_type
346+
else:
347+
return_type = AnyType(TypeOfAny.unannotated)
348+
349+
return arguments, return_type
350+
351+
352+
def copy_method_to_another_class(ctx: ClassDefContext, self_type: Instance,
353+
new_method_name: str, method_node: FuncDef) -> None:
354+
arguments, return_type = _prepare_new_method_arguments(method_node)
355+
add_method(ctx,
356+
new_method_name,
357+
args=arguments,
358+
return_type=return_type,
359+
self_type=self_type)

‎mypy_django_plugin/main.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from mypy.options import Options
99
from mypy.plugin import (
1010
AttributeContext, ClassDefContext, FunctionContext, MethodContext, Plugin,
11-
)
11+
DynamicClassDefContext)
1212
from mypy.types import Type as MypyType
1313

1414
import mypy_django_plugin.transformers.orm_lookups
@@ -17,6 +17,7 @@
1717
from mypy_django_plugin.transformers import (
1818
fields, forms, init_create, meta, querysets, request, settings,
1919
)
20+
from mypy_django_plugin.transformers.managers import create_new_manager_class_from_from_queryset_method
2021
from mypy_django_plugin.transformers.models import process_model_class
2122

2223

@@ -242,6 +243,15 @@ def get_attribute_hook(self, fullname: str
242243
return partial(request.set_auth_user_model_as_type_for_request_user, django_context=self.django_context)
243244
return None
244245

246+
def get_dynamic_class_hook(self, fullname: str
247+
) -> Optional[Callable[[DynamicClassDefContext], None]]:
248+
if fullname.endswith('from_queryset'):
249+
class_name, _, _ = fullname.rpartition('.')
250+
info = self._get_typeinfo_or_none(class_name)
251+
if info and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME):
252+
return partial(create_new_manager_class_from_from_queryset_method,
253+
django_context=self.django_context)
254+
245255

246256
def plugin(version):
247257
return NewSemanalDjangoPlugin

‎mypy_django_plugin/transformers/managers.py

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
from mypy.nodes import MemberExpr, NameExpr, TypeInfo, FuncDef
2+
from mypy.plugin import DynamicClassDefContext, ClassDefContext
3+
from mypy.types import Instance, TypeOfAny, AnyType
4+
5+
from mypy_django_plugin.django.context import DjangoContext
6+
from mypy_django_plugin.lib import helpers
7+
8+
9+
def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefContext,
10+
django_context: DjangoContext) -> None:
11+
api = helpers.get_semanal_api(ctx)
12+
current_module = api.cur_mod_node
13+
# remove variable with name
14+
del current_module.names[ctx.name]
15+
16+
assert isinstance(ctx.call.callee, MemberExpr)
17+
assert isinstance(ctx.call.callee.expr, NameExpr)
18+
base_manager_info = ctx.call.callee.expr.node
19+
if base_manager_info is None:
20+
if not api.final_iteration:
21+
api.defer()
22+
return
23+
24+
new_manager_info = helpers.add_new_class_for_module(current_module,
25+
name=ctx.name,
26+
bases=[Instance(base_manager_info,
27+
[AnyType(TypeOfAny.unannotated)])])
28+
29+
derived_queryset_fullname = ctx.call.args[0].fullname
30+
sym = api.lookup_fully_qualified_or_none(derived_queryset_fullname)
31+
assert sym is not None
32+
if sym.node is None:
33+
if not api.final_iteration:
34+
api.defer()
35+
return
36+
else:
37+
# TODO: inherit from Any to prevent false-positives, if queryset class cannot be resolved
38+
return
39+
40+
derived_queryset_info = sym.node
41+
assert isinstance(derived_queryset_info, TypeInfo)
42+
43+
class_def_context = ClassDefContext(cls=new_manager_info.defn,
44+
reason=ctx.call, api=api)
45+
self_type = Instance(new_manager_info, [])
46+
for name, sym in derived_queryset_info.names.items():
47+
if isinstance(sym.node, FuncDef):
48+
helpers.copy_method_to_another_class(class_def_context,
49+
self_type,
50+
new_method_name=name,
51+
method_node=sym.node)

‎mypy_django_plugin/transformers/models.py

Lines changed: 18 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,15 @@
22
from typing import List, Tuple, Type
33

44
from django.db.models.base import Model
5+
from django.db.models.query import QuerySet
6+
57
from django.db.models.fields import DateField, DateTimeField
68
from django.db.models.fields.related import ForeignKey
79
from django.db.models.fields.reverse_related import (
810
ManyToManyRel, ManyToOneRel, OneToOneRel,
911
)
1012
from mypy.nodes import ARG_STAR2, Argument, Context, FuncDef, TypeInfo, Var
11-
from mypy.plugin import ClassDefContext
13+
from mypy.plugin import ClassDefContext, SemanticAnalyzerPluginInterface
1214
from mypy.plugins import common
1315
from mypy.plugins.common import add_method
1416
from mypy.types import AnyType, CallableType, Instance
@@ -23,7 +25,7 @@
2325

2426
class ModelClassInitializer:
2527
def __init__(self, ctx: ClassDefContext, django_context: DjangoContext):
26-
self.api = ctx.api
28+
self.api = ctx.api # type: SemanticAnalyzerPluginInterface
2729
self.model_classdef = ctx.cls
2830
self.django_context = django_context
2931
self.ctx = ctx
@@ -132,7 +134,16 @@ def _is_manager_any(self, typ: Instance) -> bool:
132134
def run_with_model_cls(self, model_cls: Type[Model]) -> None:
133135
for manager_name, manager in model_cls._meta.managers_map.items():
134136
manager_fullname = helpers.get_class_fullname(manager.__class__)
135-
manager_info = self.lookup_typeinfo_or_incomplete_defn_error(manager_fullname)
137+
try:
138+
manager_info = self.lookup_typeinfo_or_incomplete_defn_error(manager_fullname)
139+
except helpers.IncompleteDefnException as e:
140+
if manager._queryset_class == QuerySet:
141+
raise e
142+
# manager created with from_queryset()
143+
# helpers.get_semanal_api(self.ctx).()
144+
# self.api.ana
145+
print()
146+
#
136147

137148
if manager_name not in self.model_classdef.info.names:
138149
manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])])
@@ -167,12 +178,10 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
167178
for name, sym in manager_info.names.items():
168179
# replace self type with new class, if copying method
169180
if isinstance(sym.node, FuncDef):
170-
arguments, return_type = self.prepare_new_method_arguments(sym.node)
171-
add_method(new_cls_def_context,
172-
name,
173-
args=arguments,
174-
return_type=return_type,
175-
self_type=custom_manager_type)
181+
helpers.copy_method_to_another_class(new_cls_def_context,
182+
self_type=custom_manager_type,
183+
new_method_name=name,
184+
method_node=sym.node)
176185
continue
177186

178187
new_sym = sym.copy()
@@ -185,20 +194,6 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
185194

186195
self.add_new_node_to_model_class(manager_name, custom_manager_type)
187196

188-
def prepare_new_method_arguments(self, node: FuncDef) -> Tuple[List[Argument], MypyType]:
189-
arguments = []
190-
for argument in node.arguments[1:]:
191-
if argument.type_annotation is None:
192-
argument.type_annotation = AnyType(TypeOfAny.unannotated)
193-
arguments.append(argument)
194-
195-
if isinstance(node.type, CallableType):
196-
return_type = node.type.ret_type
197-
else:
198-
return_type = AnyType(TypeOfAny.unannotated)
199-
200-
return arguments, return_type
201-
202197

203198
class AddDefaultManagerAttribute(ModelClassInitializer):
204199
def run_with_model_cls(self, model_cls: Type[Model]) -> None:

‎test-data/typecheck/managers/querysets/test_from_queryset.yml

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
- case: test_from_queryset_returns_intersection_of_manager_and_queryset
2+
main: |
3+
from myapp.models import MyModel, NewManager
4+
reveal_type(NewManager()) # N: Revealed type is 'myapp.models.NewManager'
5+
reveal_type(MyModel.objects) # N: Revealed type is 'myapp.models.NewManager'
6+
reveal_type(MyModel.objects.get()) # N: Revealed type is 'myapp.models.MyModel'
7+
reveal_type(MyModel.objects.manager_only_method()) # N: Revealed type is 'builtins.int'
8+
reveal_type(MyModel.objects.manager_and_queryset_method()) # N: Revealed type is 'builtins.str'
9+
installed_apps:
10+
- myapp
11+
files:
12+
- path: myapp/__init__.py
13+
- path: myapp/models.py
14+
content: |
15+
from django.db import models
16+
class ModelBaseManager(models.Manager):
17+
def manager_only_method(self) -> int:
18+
return 1
19+
class ModelQuerySet(models.QuerySet):
20+
def manager_and_queryset_method(self) -> str:
21+
return 'hello'
22+
23+
NewManager = ModelBaseManager.from_queryset(ModelQuerySet)
24+
class MyModel(models.Model):
25+
objects = NewManager()
26+

0 commit comments

Comments
 (0)