diff --git a/mypy_django_plugin/plugins/related_fields.py b/mypy_django_plugin/plugins/related_fields.py index a62eea6..41acb03 100644 --- a/mypy_django_plugin/plugins/related_fields.py +++ b/mypy_django_plugin/plugins/related_fields.py @@ -1,11 +1,11 @@ from typing import Optional, Callable -from mypy.nodes import SymbolTableNode, MDEF, Var +from mypy.nodes import Var, MDEF, SymbolTableNode from mypy.plugin import Plugin, FunctionContext -from mypy.types import Type, CallableType, TypeOfAny, AnyType, Instance +from mypy.types import Type, CallableType, Instance -def set_related_fields(ctx: FunctionContext) -> Type: +def set_related_name_manager_for_foreign_key(ctx: FunctionContext) -> Type: if 'related_name' not in ctx.context.arg_names: return ctx.default_return_type @@ -23,8 +23,8 @@ def set_related_fields(ctx: FunctionContext) -> Type: related_var = Var(related_name, queryset_type) related_var.info = queryset_type.type - - referred_to.type.names[related_name] = SymbolTableNode(MDEF, related_var) + referred_to.type.names[related_name] = SymbolTableNode(MDEF, related_var, + plugin_generated=True) return ctx.default_return_type @@ -32,7 +32,7 @@ class RelatedFieldsPlugin(Plugin): def get_function_hook(self, fullname: str ) -> Optional[Callable[[FunctionContext], Type]]: if fullname == 'django.db.models.fields.related.ForeignKey': - return set_related_fields + return set_related_name_manager_for_foreign_key return None diff --git a/test/plugins.ini b/test/plugins.ini index 353143d..efdce26 100644 --- a/test/plugins.ini +++ b/test/plugins.ini @@ -1,4 +1,4 @@ [mypy] plugins = mypy_django_plugin.plugins.postgres_fields, - mypy_django_plugin.plugins.related_fields \ No newline at end of file + mypy_django_plugin.plugins.related_fields diff --git a/test/test-data/check-postgres-fields.test b/test/test-data/check-postgres-fields.test index 2de1691..3dfcb07 100644 --- a/test/test-data/check-postgres-fields.test +++ b/test/test-data/check-postgres-fields.test @@ -23,6 +23,4 @@ class User(models.Model): user = User() reveal_type(user.members) # E: Revealed type is 'builtins.list[builtins.int*]' reveal_type(user.members_as_text) # E: Revealed type is 'builtins.list[builtins.str*]' -[out] - - +[out] \ No newline at end of file diff --git a/test/testdjango.py b/test/testdjango.py index d031bb2..c148ff2 100644 --- a/test/testdjango.py +++ b/test/testdjango.py @@ -16,7 +16,7 @@ class DjangoTestSuite(DataSuite): files = [ 'check-model-fields.test', 'check-postgres-fields.test', - 'check-model-relations.test' + 'check-model-relations.test', ] data_prefix = str(TEST_DATA_DIR)