Reorganize code a bit, add current directory to sys.path (#198)

mkurnikov · web-flow · commit 717be5940f8e · 2019-10-05T19:44:29.000+03:00
* reorganize code a bit

* add current directory to sys.path

* remove PYTHONPATH mention from the docs

* linting
diff --git a/README.md b/README.md
@@ -52,12 +52,6 @@ django_settings_module = mysettings
 
 Where `mysettings` is a value of `DJANGO_SETTINGS_MODULE` (with or without quotes)
 
-You might also need to explicitly tweak your `PYTHONPATH` the very same way `django` does it internally in case you have troubles with mypy / django plugin not finding your settings module. Try adding the root path of your project to your `PYTHONPATH` environment variable like so:
-
-```bash
-PYTHONPATH=${PYTHONPATH}:${PWD}
-```
-
 Current implementation uses Django runtime to extract models information, so it will crash, if your installed apps `models.py` is not correct. For this same reason, you cannot use `reveal_type` inside global scope of any Python file that will be executed for `django.setup()`. 
 
 In other words, if your `manage.py runserver` crashes, mypy will crash too. 
diff --git a/mypy_django_plugin/django/context.py b/mypy_django_plugin/django/context.py
diff --git a/mypy_django_plugin/lib/fullnames.py b/mypy_django_plugin/lib/fullnames.py
@@ -37,3 +37,5 @@
 MIGRATION_CLASS_FULLNAME = 'django.db.migrations.migration.Migration'
 OPTIONS_CLASS_FULLNAME = 'django.db.models.options.Options'
 HTTPREQUEST_CLASS_FULLNAME = 'django.http.request.HttpRequest'
+
+F_EXPRESSION_FULLNAME = 'django.db.models.expressions.F'
diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py
@@ -284,3 +284,11 @@ def get_typechecker_api(ctx: Union[AttributeContext, MethodContext, FunctionCont
 def is_model_subclass_info(info: TypeInfo, django_context: 'DjangoContext') -> bool:
     return (info.fullname() in django_context.model_base_classes
             or info.has_base(fullnames.MODEL_CLASS_FULLNAME))
+
+
+def check_types_compatible(ctx: Union[FunctionContext, MethodContext],
+                           *, expected_type: MypyType, actual_type: MypyType, error_message: str) -> None:
+    api = get_typechecker_api(ctx)
+    api.check_subtype(actual_type, expected_type,
+                      ctx.context, error_message,
+                      'got', 'expected')
diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py
@@ -11,6 +11,7 @@
 )
 from mypy.types import Type as MypyType
 
+import mypy_django_plugin.transformers.orm_lookups
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.lib import fullnames, helpers
 from mypy_django_plugin.transformers import (
@@ -148,13 +149,13 @@ def get_additional_deps(self, file: MypyFile) -> List[Tuple[int, str, int]]:
             # forward relations
             for field in self.django_context.get_model_fields(model_class):
                 if isinstance(field, RelatedField):
-                    related_model_cls = self.django_context.fields_context.get_related_model_cls(field)
+                    related_model_cls = self.django_context.get_field_related_model_cls(field)
                     related_model_module = related_model_cls.__module__
                     if related_model_module != file.fullname():
                         deps.add(self._new_dependency(related_model_module))
             # reverse relations
             for relation in model_class._meta.related_objects:
-                related_model_cls = self.django_context.fields_context.get_related_model_cls(relation)
+                related_model_cls = self.django_context.get_field_related_model_cls(relation)
                 related_model_module = related_model_cls.__module__
                 if related_model_module != file.fullname():
                     deps.add(self._new_dependency(related_model_module))
@@ -210,7 +211,8 @@ def get_method_hook(self, fullname: str
         if class_fullname in manager_classes and method_name == 'create':
             return partial(init_create.redefine_and_typecheck_model_create, django_context=self.django_context)
         if class_fullname in manager_classes and method_name in {'filter', 'get', 'exclude'}:
-            return partial(init_create.typecheck_queryset_filter, django_context=self.django_context)
+            return partial(mypy_django_plugin.transformers.orm_lookups.typecheck_queryset_filter,
+                           django_context=self.django_context)
         return None
 
     def get_base_class_hook(self, fullname: str
diff --git a/mypy_django_plugin/transformers/fields.py b/mypy_django_plugin/transformers/fields.py
@@ -44,7 +44,7 @@ def fill_descriptor_types_for_related_field(ctx: FunctionContext, django_context
 
     assert isinstance(current_field, RelatedField)
 
-    related_model_cls = django_context.fields_context.get_related_model_cls(current_field)
+    related_model_cls = django_context.get_field_related_model_cls(current_field)
 
     related_model = related_model_cls
     related_model_to_set = related_model_cls
diff --git a/mypy_django_plugin/transformers/init_create.py b/mypy_django_plugin/transformers/init_create.py
@@ -6,7 +6,7 @@
 from mypy.types import Type as MypyType
 
 from mypy_django_plugin.django.context import DjangoContext
-from mypy_django_plugin.lib import fullnames, helpers
+from mypy_django_plugin.lib import helpers
 
 
 def get_actual_types(ctx: Union[MethodContext, FunctionContext],
@@ -30,12 +30,6 @@ def get_actual_types(ctx: Union[MethodContext, FunctionContext],
     return actual_types
 
 
-def check_types_compatible(ctx, *, expected_type, actual_type, error_message):
-    ctx.api.check_subtype(actual_type, expected_type,
-                          ctx.context, error_message,
-                          'got', 'expected')
-
-
 def typecheck_model_method(ctx: Union[FunctionContext, MethodContext], django_context: DjangoContext,
                            model_cls: Type[Model], method: str) -> MypyType:
     typechecker_api = helpers.get_typechecker_api(ctx)
@@ -48,11 +42,11 @@ def typecheck_model_method(ctx: Union[FunctionContext, MethodContext], django_co
                                                                            model_cls.__name__),
                          ctx.context)
             continue
-        check_types_compatible(ctx,
-                               expected_type=expected_types[actual_name],
-                               actual_type=actual_type,
-                               error_message='Incompatible type for "{}" of "{}"'.format(actual_name,
-                                                                                         model_cls.__name__))
+        helpers.check_types_compatible(ctx,
+                                       expected_type=expected_types[actual_name],
+                                       actual_type=actual_type,
+                                       error_message='Incompatible type for "{}" of "{}"'.format(actual_name,
+                                                                                                 model_cls.__name__))
 
     return ctx.default_return_type
 
@@ -79,40 +73,3 @@ def redefine_and_typecheck_model_create(ctx: MethodContext, django_context: Djan
         return ctx.default_return_type
 
     return typecheck_model_method(ctx, django_context, model_cls, 'create')
-
-
-def typecheck_queryset_filter(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
-    lookup_kwargs = ctx.arg_names[1]
-    provided_lookup_types = ctx.arg_types[1]
-
-    assert isinstance(ctx.type, Instance)
-
-    if not ctx.type.args or not isinstance(ctx.type.args[0], Instance):
-        return ctx.default_return_type
-
-    model_cls_fullname = ctx.type.args[0].type.fullname()
-    model_cls = django_context.get_model_class_by_fullname(model_cls_fullname)
-    if model_cls is None:
-        return ctx.default_return_type
-
-    for lookup_kwarg, provided_type in zip(lookup_kwargs, provided_lookup_types):
-        if lookup_kwarg is None:
-            continue
-        # Combinables are not supported yet
-        if (isinstance(provided_type, Instance)
-                and provided_type.type.has_base('django.db.models.expressions.Combinable')):
-            continue
-
-        lookup_type = django_context.lookups_context.resolve_lookup_expected_type(ctx, model_cls, lookup_kwarg)
-        # Managers as provided_type is not supported yet
-        if (isinstance(provided_type, Instance)
-                and helpers.has_any_of_bases(provided_type.type, (fullnames.MANAGER_CLASS_FULLNAME,
-                                                                  fullnames.QUERYSET_CLASS_FULLNAME))):
-            return ctx.default_return_type
-
-        check_types_compatible(ctx,
-                               expected_type=lookup_type,
-                               actual_type=provided_type,
-                               error_message=f'Incompatible type for lookup {lookup_kwarg!r}:')
-
-    return ctx.default_return_type
diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py
@@ -99,10 +99,10 @@ class AddRelatedModelsId(ModelClassInitializer):
     def run_with_model_cls(self, model_cls: Type[Model]) -> None:
         for field in model_cls._meta.get_fields():
             if isinstance(field, ForeignKey):
-                related_model_cls = self.django_context.fields_context.get_related_model_cls(field)
+                related_model_cls = self.django_context.get_field_related_model_cls(field)
                 rel_primary_key_field = self.django_context.get_primary_key_field(related_model_cls)
                 field_info = self.lookup_class_typeinfo_or_incomplete_defn_error(rel_primary_key_field.__class__)
-                is_nullable = self.django_context.fields_context.get_field_nullability(field, None)
+                is_nullable = self.django_context.get_field_nullability(field, None)
                 set_type, get_type = get_field_descriptor_types(field_info, is_nullable)
                 self.add_new_node_to_model_class(field.attname,
                                                  Instance(field_info, [set_type, get_type]))
@@ -162,7 +162,7 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
                 # no reverse accessor
                 continue
 
-            related_model_cls = self.django_context.fields_context.get_related_model_cls(relation)
+            related_model_cls = self.django_context.get_field_related_model_cls(relation)
             related_model_info = self.lookup_class_typeinfo_or_incomplete_defn_error(related_model_cls)
 
             if isinstance(relation, OneToOneRel):
diff --git a/mypy_django_plugin/transformers/orm_lookups.py b/mypy_django_plugin/transformers/orm_lookups.py
@@ -0,0 +1,51 @@
+from mypy.plugin import MethodContext
+from mypy.types import AnyType, Instance
+from mypy.types import Type as MypyType
+from mypy.types import TypeOfAny
+
+from mypy_django_plugin.django.context import DjangoContext
+from mypy_django_plugin.lib import fullnames, helpers
+
+
+def typecheck_queryset_filter(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
+    lookup_kwargs = ctx.arg_names[1]
+    provided_lookup_types = ctx.arg_types[1]
+
+    assert isinstance(ctx.type, Instance)
+
+    if not ctx.type.args or not isinstance(ctx.type.args[0], Instance):
+        return ctx.default_return_type
+
+    model_cls_fullname = ctx.type.args[0].type.fullname()
+    model_cls = django_context.get_model_class_by_fullname(model_cls_fullname)
+    if model_cls is None:
+        return ctx.default_return_type
+
+    for lookup_kwarg, provided_type in zip(lookup_kwargs, provided_lookup_types):
+        if lookup_kwarg is None:
+            continue
+        if (isinstance(provided_type, Instance)
+                and provided_type.type.has_base('django.db.models.expressions.Combinable')):
+            provided_type = resolve_combinable_type(provided_type, django_context)
+
+        lookup_type = django_context.resolve_lookup_expected_type(ctx, model_cls, lookup_kwarg)
+        # Managers as provided_type is not supported yet
+        if (isinstance(provided_type, Instance)
+                and helpers.has_any_of_bases(provided_type.type, (fullnames.MANAGER_CLASS_FULLNAME,
+                                                                  fullnames.QUERYSET_CLASS_FULLNAME))):
+            return ctx.default_return_type
+
+        helpers.check_types_compatible(ctx,
+                                       expected_type=lookup_type,
+                                       actual_type=provided_type,
+                                       error_message=f'Incompatible type for lookup {lookup_kwarg!r}:')
+
+    return ctx.default_return_type
+
+
+def resolve_combinable_type(combinable_type: Instance, django_context: DjangoContext) -> MypyType:
+    if combinable_type.type.fullname() != fullnames.F_EXPRESSION_FULLNAME:
+        # Combinables aside from F expressions are unsupported
+        return AnyType(TypeOfAny.explicit)
+
+    return django_context.resolve_f_expression_type(combinable_type)
diff --git a/mypy_django_plugin/transformers/querysets.py b/mypy_django_plugin/transformers/querysets.py
@@ -40,19 +40,19 @@ def determine_proper_manager_type(ctx: FunctionContext) -> MypyType:
 def get_field_type_from_lookup(ctx: MethodContext, django_context: DjangoContext, model_cls: Type[Model],
                                *, method: str, lookup: str) -> Optional[MypyType]:
     try:
-        lookup_field = django_context.lookups_context.resolve_lookup_info_field(model_cls, lookup)
+        lookup_field = django_context.resolve_lookup_into_field(model_cls, lookup)
     except FieldError as exc:
         ctx.api.fail(exc.args[0], ctx.context)
         return None
     except LookupsAreUnsupported:
         return AnyType(TypeOfAny.explicit)
 
     if isinstance(lookup_field, RelatedField) and lookup_field.column == lookup:
-        related_model_cls = django_context.fields_context.get_related_model_cls(lookup_field)
+        related_model_cls = django_context.get_field_related_model_cls(lookup_field)
         lookup_field = django_context.get_primary_key_field(related_model_cls)
 
-    field_get_type = django_context.fields_context.get_field_get_type(helpers.get_typechecker_api(ctx),
-                                                                      lookup_field, method=method)
+    field_get_type = django_context.get_field_get_type(helpers.get_typechecker_api(ctx),
+                                                       lookup_field, method=method)
     return field_get_type
 
 
@@ -73,8 +73,8 @@ def get_values_list_row_type(ctx: MethodContext, django_context: DjangoContext,
         elif named:
             column_types: 'OrderedDict[str, MypyType]' = OrderedDict()
             for field in django_context.get_model_fields(model_cls):
-                column_type = django_context.fields_context.get_field_get_type(typechecker_api, field,
-                                                                               method='values_list')
+                column_type = django_context.get_field_get_type(typechecker_api, field,
+                                                                method='values_list')
                 column_types[field.attname] = column_type
             return helpers.make_oneoff_named_tuple(typechecker_api, 'Row', column_types)
         else:
diff --git a/test-data/typecheck/managers/querysets/test_filter.yml b/test-data/typecheck/managers/querysets/test_filter.yml
@@ -212,4 +212,45 @@
                 class User(models.Model):
                     pass
                 class Profile(models.Model):
-                    user = models.OneToOneField(User, on_delete=models.CASCADE, related_name='profile')
+                    user = models.OneToOneField(User, on_delete=models.CASCADE, related_name='profile')
+
+
+# TODO
+-   case: f_expression_simple_case
+    main: |
+        from myapp.models import User
+        from django.db import models
+        User.objects.filter(username=models.F('username2'))
+        User.objects.filter(username=models.F('age'))
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+
+                class User(models.Model):
+                    username = models.TextField()
+                    username2 = models.TextField()
+
+                    age = models.IntegerField()
+
+
+# TODO
+-   case: f_expression_with_expression_math_is_not_supported
+    main: |
+        from myapp.models import User
+        from django.db import models
+        User.objects.filter(username=models.F('username2') + 'hello')
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                class User(models.Model):
+                    username = models.TextField()
+                    username2 = models.TextField()
+                    age = models.IntegerField()
