From 1a70900ba44f2497c853d873f8bcf39cf267a2d5 Mon Sep 17 00:00:00 2001
From: Andrey Maslov <greyzmeem@gmail.com>
Date: Sun, 15 Apr 2018 04:58:18 +0300
Subject: [PATCH] Fix Django 2 related object descriptor cache

---
 modeltranslation/fields.py      |  9 +++++++++
 modeltranslation/tests/tests.py | 19 ++++++++++++++++++-
 modeltranslation/translator.py  |  9 +++++++++
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/modeltranslation/fields.py b/modeltranslation/fields.py
index 9ded57b..f22350f 100644
--- a/modeltranslation/fields.py
+++ b/modeltranslation/fields.py
@@ -404,6 +404,15 @@ class LanguageCacheSingleObjectDescriptor(object):
 
     @property
     def cache_name(self):
+        """
+        Used in django 1.x
+        """
         lang = get_language()
         cache = build_localized_fieldname(self.accessor, lang)
         return "_%s_cache" % cache
+
+    def get_cache_name(self):
+        """
+        Used in django 2.x
+        """
+        return build_localized_fieldname(self.accessor, get_language())
diff --git a/modeltranslation/tests/tests.py b/modeltranslation/tests/tests.py
index afe5f1a..22d2b23 100644
--- a/modeltranslation/tests/tests.py
+++ b/modeltranslation/tests/tests.py
@@ -2960,7 +2960,8 @@ class TestManager(ModeltranslationTestBase):
         o.title = "bla"
         self.assertEqual(o.title, "bla")
 
-    def test_select_related(self):
+    @skipUnless(django.VERSION[0] == 1, 'Applicable only to django 1.x')
+    def test_select_related_django_1(self):
         test = models.TestModel.objects.create(title_de='title_de', title_en='title_en')
         with auto_populate('all'):
             models.ForeignKeyModel.objects.create(untrans=test)
@@ -2975,6 +2976,22 @@ class TestManager(ModeltranslationTestBase):
         # untrans is nullable so not included when select_related=True
         self.assertNotIn('_untrans_cache', fk_qs.select_related()[0].__dict__)
 
+    @skipUnless(django.VERSION[0] == 2, 'Applicable only to django 2.x')
+    def test_select_related_django_2(self):
+        test = models.TestModel.objects.create(title_de='title_de', title_en='title_en')
+        with auto_populate('all'):
+            models.ForeignKeyModel.objects.create(untrans=test)
+
+        fk_qs = models.ForeignKeyModel.objects.all()
+        self.assertNotIn('untrans', fk_qs[0]._state.fields_cache)
+        self.assertIn('untrans', fk_qs.select_related('untrans')[0]._state.fields_cache)
+        self.assertNotIn(
+            'untrans',
+            fk_qs.select_related('untrans').select_related(None)[0]._state.fields_cache
+        )
+        # untrans is nullable so not included when select_related=True
+        self.assertNotIn('untrans', fk_qs.select_related()[0]._state.fields_cache)
+
     def test_translation_fields_appending(self):
         from modeltranslation.manager import append_lookup_keys, append_lookup_key
         self.assertEqual(set(['untrans']), append_lookup_key(models.ForeignKeyModel, 'untrans'))
diff --git a/modeltranslation/translator.py b/modeltranslation/translator.py
index 3a9a297..962d25d 100644
--- a/modeltranslation/translator.py
+++ b/modeltranslation/translator.py
@@ -1,4 +1,6 @@
 # -*- coding: utf-8 -*-
+from functools import partial
+
 from django import VERSION
 from django.utils.six import with_metaclass
 from django.core.exceptions import ImproperlyConfigured
@@ -428,6 +430,13 @@ def patch_related_object_descriptor_caching(ro_descriptor):
     """
     class NewSingleObjectDescriptor(LanguageCacheSingleObjectDescriptor, ro_descriptor.__class__):
         pass
+
+    if VERSION[0] == 2:
+        ro_descriptor.related.get_cache_name = partial(
+            NewSingleObjectDescriptor.get_cache_name,
+            ro_descriptor,
+        )
+
     ro_descriptor.accessor = ro_descriptor.related.get_accessor_name()
     ro_descriptor.__class__ = NewSingleObjectDescriptor