Skip to content

Commit

Permalink
Always create manager at module level
Browse files Browse the repository at this point in the history
When the manager is added at the class level, which happened when it was
created inline in the model body, it's not possible to retrieve the
manager again based on fullname. That lead to problems with inheritance
and the default manager.
  • Loading branch information
ljodal committed Jul 3, 2022
1 parent aa86858 commit a0d8d6d
Show file tree
Hide file tree
Showing 4 changed files with 56 additions and 33 deletions.
28 changes: 18 additions & 10 deletions mypy_django_plugin/lib/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -208,6 +208,23 @@ def is_annotated_model_fullname(model_cls_fullname: str) -> bool:
return model_cls_fullname.startswith(WITH_ANNOTATIONS_FULLNAME + "[")


def create_type_info(name: str, module: str, bases: list[Instance]) -> TypeInfo:

# make new class expression
classdef = ClassDef(name, Block([]))
classdef.fullname = module + "." + name

# make new TypeInfo
new_typeinfo = TypeInfo(SymbolTable(), classdef, module)
new_typeinfo.bases = bases
calculate_mro(new_typeinfo)
new_typeinfo.calculate_metaclass_type()

classdef.info = new_typeinfo

return new_typeinfo


def add_new_class_for_module(
module: MypyFile,
name: str,
Expand All @@ -217,15 +234,7 @@ def add_new_class_for_module(
) -> TypeInfo:
new_class_unique_name = checker.gen_unique_name(name, module.names)

# make new class expression
classdef = ClassDef(new_class_unique_name, Block([]))
classdef.fullname = module.fullname + "." + new_class_unique_name

# make new TypeInfo
new_typeinfo = TypeInfo(SymbolTable(), classdef, module.fullname)
new_typeinfo.bases = bases
calculate_mro(new_typeinfo)
new_typeinfo.calculate_metaclass_type()
new_typeinfo = create_type_info(new_class_unique_name, module.fullname, bases)

# add fields
if fields:
Expand All @@ -237,7 +246,6 @@ def add_new_class_for_module(
MDEF, var, plugin_generated=True, no_serialize=no_serialize
)

classdef.info = new_typeinfo
module.names[new_class_unique_name] = SymbolTableNode(
GDEF, new_typeinfo, plugin_generated=True, no_serialize=no_serialize
)
Expand Down
14 changes: 7 additions & 7 deletions mypy_django_plugin/transformers/managers.py
Original file line number Diff line number Diff line change
Expand Up @@ -199,12 +199,6 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
# So that the plugin will reparameterize the manager when it is constructed inside of a Model definition
helpers.add_new_manager_base(semanal_api, new_manager_info.fullname)

# Insert the new manager (dynamic) class
assert semanal_api.add_symbol_table_node(
ctx.name,
SymbolTableNode(GDEF, new_manager_info, plugin_generated=True),
)


def create_manager_info_from_from_queryset_call(
api: SemanticAnalyzerPluginInterface, call_expr: CallExpr, name: Optional[str] = None
Expand Down Expand Up @@ -251,6 +245,12 @@ def create_manager_info_from_from_queryset_call(
base_manager_info.metadata.setdefault("from_queryset_managers", {})
base_manager_info.metadata["from_queryset_managers"][manager_fullname] = new_manager_info.fullname

# Add the new manager to the current module
module = api.modules[api.cur_mod_id]
module.names[name or manager_name] = SymbolTableNode(
GDEF, new_manager_info, plugin_generated=True, no_serialize=False
)

return new_manager_info


Expand All @@ -261,7 +261,7 @@ def create_manager_class(
base_manager_instance = fill_typevars(base_manager_info)
assert isinstance(base_manager_instance, Instance)

manager_info = api.basic_new_typeinfo(name, basetype_or_fallback=base_manager_instance, line=line)
manager_info = helpers.create_type_info(name, api.cur_mod_id, bases=[base_manager_instance])
manager_info.line = line
manager_info.type_vars = base_manager_info.type_vars
manager_info.defn.type_vars = base_manager_info.defn.type_vars
Expand Down
16 changes: 2 additions & 14 deletions mypy_django_plugin/transformers/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@
from mypy_django_plugin.errorcodes import MANAGER_MISSING
from mypy_django_plugin.lib import fullnames, helpers
from mypy_django_plugin.lib.fullnames import ANNOTATIONS_FULLNAME, ANY_ATTR_ALLOWED_CLASS_FULLNAME, MODEL_CLASS_FULLNAME
from mypy_django_plugin.lib.helpers import add_new_class_for_module
from mypy_django_plugin.transformers import fields
from mypy_django_plugin.transformers.fields import get_field_descriptor_types
from mypy_django_plugin.transformers.managers import create_manager_info_from_from_queryset_call
Expand Down Expand Up @@ -357,18 +356,7 @@ class MyModel(models.Model):
if not isinstance(expr, CallExpr) or not isinstance(expr.callee, CallExpr):
return None

new_manager_info = create_manager_info_from_from_queryset_call(self.api, expr.callee)

if new_manager_info:
assert self.api.add_symbol_table_node(
# We'll use `new_manager_info.name` instead of `manager_class_name` here
# to handle possible name collisions, as it's unique.
new_manager_info.name,
# Note that the generated manager type is always inserted at module level
SymbolTableNode(MDEF, new_manager_info, plugin_generated=True),
)

return new_manager_info
return create_manager_info_from_from_queryset_call(self.api, expr.callee)


class AddDefaultManagerAttribute(ModelClassInitializer):
Expand Down Expand Up @@ -646,7 +634,7 @@ def get_or_create_annotated_type(
else:
annotated_model_type = api.named_generic_type(ANY_ATTR_ALLOWED_CLASS_FULLNAME, [])

annotated_typeinfo = add_new_class_for_module(
annotated_typeinfo = helpers.add_new_class_for_module(
model_module_file,
type_name,
bases=[model_type] if fields_dict is not None else [model_type, annotated_model_type],
Expand Down
31 changes: 29 additions & 2 deletions tests/typecheck/managers/querysets/test_from_queryset.yml
Original file line number Diff line number Diff line change
Expand Up @@ -362,12 +362,14 @@
- case: test_queryset_in_model_class_body
main: |
from myapp.models import MyModel
reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.MyModel.MyManagerFromMyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.all) # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.custom) # N: Revealed type is "def () -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.all().filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.custom().filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects2) # N: Revealed type is "myapp.models.MyModel.MyManagerFromMyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects2) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]"
installed_apps:
- myapp
files:
Expand All @@ -387,6 +389,31 @@
objects = MyManager.from_queryset(MyQuerySet)()
objects2 = MyManager.from_queryset(MyQuerySet)()
- case: test_queryset_in_model_class_body_subclass
main: |
from myapp.models import MyModel
reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.BaseManagerFromBaseQuerySet[myapp.models.MyModel]"
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
class BaseManager(models.Manager["BaseModel"]):
pass
class BaseQuerySet(models.QuerySet["BaseModel"]):
def custom(self) -> "BaseQuerySet":
pass
class BaseModel(models.Model):
objects = BaseManager.from_queryset(BaseQuerySet)()
class MyModel(BaseModel):
pass
- case: from_queryset_includes_methods_returning_queryset
main: |
from myapp.models import MyModel
Expand Down

0 comments on commit a0d8d6d

Please sign in to comment.