Repository: adamchainz/django-mysql

Issue:
Make QuerySetMixin capable of wrapping an existing QuerySet
It would be nice to be able to use the extensions with models defined in other apps, without much hackery. For example I want to do something like:

``` py
from django.contrib.admin.models import LogEntry
qs = LogEntry.objects.filter(user=x)
qs = QuerySetMixin.wrap(qs).force_index('bla')
```

N.B. old versions of Django used to make this a bit easier in an undocumented fashion - `_clone` took a `klass` parameter that allowed mutation, so a dynamic subclass could be generated. We could do something similar. However wrapped `QuerySet`s are at a risk of being un-pickle-able if the class is purely dynamic, I know this needs some extra stuff to make it work with pickle.



Discussion:

Resolving patch:
diff --git a/django_mysql/models/__init__.py b/django_mysql/models/__init__.py
--- a/django_mysql/models/__init__.py
+++ b/django_mysql/models/__init__.py
@@ -4,8 +4,8 @@
 )
 from django_mysql.models.expressions import ListF, SetF  # noqa
 from django_mysql.models.query import (  # noqa
-    ApproximateInt, SmartChunkedIterator, SmartIterator, pt_visual_explain,
-    QuerySet, QuerySetMixin
+    add_QuerySetMixin, ApproximateInt, SmartChunkedIterator, SmartIterator,
+    pt_visual_explain, QuerySet, QuerySetMixin
 )
 from django_mysql.models.fields import (  # noqa
     Bit1BooleanField, DynamicField, ListCharField, ListTextField,
diff --git a/django_mysql/models/query.py b/django_mysql/models/query.py
--- a/django_mysql/models/query.py
+++ b/django_mysql/models/query.py
@@ -31,7 +31,11 @@ def __init__(self, *args, **kwargs):
 
     def _clone(self, *args, **kwargs):
         clone = super(QuerySetMixin, self)._clone(*args, **kwargs)
-        clone._count_tries_approx = copy(self._count_tries_approx)
+
+        clone._count_tries_approx = copy(
+            getattr(self, '_count_tries_approx', False)
+        )
+
         if hasattr(self, '_found_rows'):
             # If it's a number, don't copy it - the clone has a fresh result
             # cache
@@ -217,6 +221,30 @@ class QuerySet(QuerySetMixin, models.QuerySet):
     pass
 
 
+def add_QuerySetMixin(queryset):
+    queryset2 = queryset._clone()
+    queryset2.__class__ = _make_mixin_class(queryset.__class__)
+    return queryset2
+
+
+_mixin_classes = {}
+
+
+def _make_mixin_class(klass):
+    global _mixin_classes
+
+    if klass not in _mixin_classes:
+        class MixedInQuerySet(QuerySetMixin, klass):
+            pass
+
+        if six.PY2:
+            MixedInQuerySet.__name__ = b'MySQL' + klass.__name__
+        else:
+            MixedInQuerySet.__name__ = 'MySQL' + klass.__name__
+        _mixin_classes[klass] = MixedInQuerySet
+    return _mixin_classes[klass]
+
+
 @six.python_2_unicode_compatible
 class ApproximateInt(int):
     """


Test patch:
diff --git a/tests/testapp/test_models.py b/tests/testapp/test_models.py
--- a/tests/testapp/test_models.py
+++ b/tests/testapp/test_models.py
@@ -5,12 +5,15 @@
 from unittest import skipUnless
 
 import pytest
+from django.contrib.auth.models import User
 from django.db.models.query import QuerySet
 from django.template import Context, Template
 from django.test import TestCase
 from django.test.utils import override_settings
 
-from django_mysql.models import ApproximateInt, SmartIterator
+from django_mysql.models import (
+    ApproximateInt, SmartIterator, add_QuerySetMixin
+)
 from django_mysql.utils import have_program, index_name
 from testapp.models import (
     Author, AuthorExtra, AuthorMultiIndex, Book, NameAuthor, VanillaAuthor
@@ -18,6 +21,18 @@
 from testapp.utils import CaptureLastQuery, captured_stdout, used_indexes
 
 
+class MixinQuerysetTests(TestCase):
+
+    def test_User(self):
+        qs = User.objects.all()
+        assert not hasattr(qs, 'approx_count')
+        qs2 = add_QuerySetMixin(qs)
+
+        assert hasattr(qs2, 'approx_count')
+        assert qs.__class__ != qs2.__class__
+        assert qs.__class__ in qs2.__class__.__mro__
+
+
 class ApproximateCountTests(TestCase):
 
     def setUp(self):