From 65c4cd081617299fde4780eb110f507428e2b128 Mon Sep 17 00:00:00 2001
From: James Outterside <james.outterside@gmail.com>
Date: Tue, 24 Mar 2020 13:09:01 +0000
Subject: [PATCH] Added manager to module when casting so importable via
 migrations, removes need for deconstruct.

---
 gdpr_assist/models.py    | 39 +++++++++++++++++----------------------
 tests/test_model_meta.py |  4 ++--
 2 files changed, 19 insertions(+), 24 deletions(-)

diff --git a/gdpr_assist/models.py b/gdpr_assist/models.py
index 48f546b..a88c7a9 100644
--- a/gdpr_assist/models.py
+++ b/gdpr_assist/models.py
@@ -3,6 +3,7 @@
 """
 from copy import copy
 import six
+import sys
 
 from django.apps import apps
 from django.db import models
@@ -47,9 +48,13 @@ def _cast_class(cls, queryset):
         """
         # Make a subclass of PrivacyQuerySet and the original class
         orig_cls = queryset.__class__
-        queryset.__class__ = type(
-            str('CastPrivacy{}'.format(orig_cls.__name__)), (cls, orig_cls), {},
-        )
+        new_cls_name = str('CastPrivacy{}'.format(orig_cls.__name__))
+        queryset.__class__ = type(new_cls_name, (cls, orig_cls), {},)
+
+        # add to current module
+        current_module = sys.modules[__name__]
+        setattr(current_module, new_cls_name, queryset.__class__)
+
         return queryset
 
 
@@ -85,29 +90,19 @@ def _cast_class(cls, manager):
         The new class is given the same name as the old class, but with the prefix
         'CastPrivacy' to indicate the type of the object has changed, eg a normal
         Manager will become CastPrivacyManager
+
+        Also add the new manager to the module, so it can be imported for migrations.
         """
         # Make a subclass of PrivacyQuerySet and the original class
         orig_cls = manager.__class__
-        manager.__class__ = type(
-            str('CastPrivacy{}'.format(orig_cls.__name__)), (cls, orig_cls), {},
-        )
-        return manager
+        new_cls_name = str('CastPrivacy{}'.format(orig_cls.__name__))
+        manager.__class__ = type(new_cls_name, (cls, orig_cls), {},)
 
-    def deconstruct(self):
-        """
-        Deconstruct the original manager - it will be cast again next time.
-        """
-        # Check bases are as expected from _cast_class
-        bases = self.__class__.__bases__
-        if len(bases) != 2:  # pragma: no cover
-            raise ValueError('Unexpected base classes for CastPrivacyManager')
-
-        # Original is second - instatiate and deconstruct it
-        orig_cls = bases[1]
-        orig_args = self._constructor_args[0]
-        orig_kwargs = self._constructor_args[1]
-        orig_manager = orig_cls(*orig_args, **orig_kwargs)
-        return orig_manager.deconstruct()
+        # add to current module
+        current_module = sys.modules[__name__]
+        setattr(current_module, new_cls_name, manager.__class__)
+
+        return manager
 
 
 class PrivacyMeta(object):
diff --git a/tests/test_model_meta.py b/tests/test_model_meta.py
index d6424ba..b3a941c 100644
--- a/tests/test_model_meta.py
+++ b/tests/test_model_meta.py
@@ -218,9 +218,9 @@ class TestRegisteredModelMigration(MigrationTestCase):
     Check registered models can be migratated
     """
     def test_manager_deconstruct__deconstructs(self):
-        # This should serialise to the original manager
+        # This should serialise to the privacy manager
         string, imports = self.serialize(ModelWithPrivacyMeta.objects)
-        self.assertEqual(string, 'django.db.models.manager.Manager()')
+        self.assertEqual(string, 'gdpr_assist.models.CastPrivacyManager()')
 
         # And check it serialises back
         obj = self.serialize_round_trip(ModelWithPrivacyMeta.objects)