Merge pull request jazzband#27 from IntuitiveWebSolutions/alexanderan…

…ikeev/-/add-ability-to-use-custom-managers

Add ability to use custom model managers that override EAV manager

eav/__init__.py

@@ -1,7 +1,7 @@
 from django.core.exceptions import ImproperlyConfigured
 from django.db.models import ForeignKey
 
-__version__ = "0.15.2"
+__version__ = "0.15.4"
 
 
 def register(model_cls, config_cls=None, container_field="container"):

eav/managers.py

@@ -45,14 +45,16 @@ def create(self, **kwargs):
         obj.save()
         return obj
 
-    def get_or_create(self, **kwargs):
+    def get_or_create(self, defaults=None, **kwargs):
         """
         Reproduces the behavior of get_or_create, eav friendly.
         """
         try:
             return self.get(**kwargs), False
         except self.model.DoesNotExist:
-            return self.create(**kwargs), True
+            if not defaults:
+                defaults = {}
+            return self.create(**defaults, **kwargs), True
 
 
 class EAVContainerManager(models.Manager):

eav/queryset.py

@@ -278,6 +278,16 @@ class EavQuerySet(QuerySet):
     Overrides relational operators for EAV models.
     """
 
+    def as_manager(cls):
+        from eav.managers import EntityManager
+
+        manager = EntityManager.from_queryset(cls)()
+        manager._built_with_as_manager = True
+        return manager
+
+    as_manager.queryset_only = True
+    as_manager = classmethod(as_manager)
+
     @eav_filter
     def filter(self, *args, **kwargs):
         """

eav/registry.py

@@ -104,15 +104,22 @@ def _attach_manager(self):
         """
         Attach the manager to *manager_attr* specified in *config_cls*
         """
-        # Save the old manager if the attribute name conflicts with the new one.
-        if hasattr(self.model_cls, self.config_cls.manager_attr):
-            mgr = getattr(self.model_cls, self.config_cls.manager_attr)
-            self.config_cls.old_mgr = mgr
-            self.model_cls._meta.local_managers.remove(mgr)
-            self.model_cls._meta._expire_cache()
+        # Save the old manager if the attribute name conflicts with the new one and old manager does not extend
+        # EntityManager
+        original_manager = getattr(self.model_cls, self.config_cls.manager_attr, None)
+        if original_manager:
+            if not isinstance(original_manager, EntityManager):
+                mgr = getattr(self.model_cls, self.config_cls.manager_attr)
+                self.config_cls.old_mgr = mgr
+                self.model_cls._meta.local_managers.remove(mgr)
+                self.model_cls._meta._expire_cache()
+                mgr = EntityManager()
+            else:
+                mgr = original_manager
+        else:
+            mgr = EntityManager()
 
         # Attach the new manager to the model.
-        mgr = EntityManager()
         mgr.contribute_to_class(self.model_cls, self.config_cls.manager_attr)
 
     def _detach_manager(self):

setup.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.15.2
+current_version = 0.15.4
 commit = True
 tag = True
 

tests/factories.py

@@ -1,7 +1,8 @@
 import factory
 
 from eav.models import Attribute, EnumValue, EnumGroup
-from .models import Container, Patient, Encounter, ExampleModel
+from .models import Container, Patient, Encounter, ExampleModel, ExampleModelWithCustomManager, \
+    ExampleModelWithCustomQuerySet
 from .utils import register
 
 
@@ -89,6 +90,24 @@ class Meta:
         model = ExampleModel
 
 
+@register
+class ExampleModelWithCustomQuerySetFactory(factory.DjangoModelFactory):
+    name = factory.Sequence(lambda n: "container-name-{}".format(n))
+    container = factory.SubFactory(ContainerFactory)
+
+    class Meta:
+        model = ExampleModelWithCustomQuerySet
+
+
+@register
+class ExampleModelWithCustomManagerFactory(factory.DjangoModelFactory):
+    name = factory.Sequence(lambda n: "container-name-{}".format(n))
+    container = factory.SubFactory(ContainerFactory)
+
+    class Meta:
+        model = ExampleModelWithCustomManager
+
+
 @register
 class PatientFactory(factory.DjangoModelFactory):
     name = factory.Sequence(lambda n: "container-name-{}".format(n))

tests/models.py

@@ -2,8 +2,13 @@
 
 from django.contrib.auth.models import AbstractUser
 from django.db import models
+from django.db.models import QuerySet
+
+from eav.managers import EntityManager
+
 from eav.decorators import register_eav
 from eav.models import AbstractEAVContainer
+from eav.queryset import EavQuerySet
 
 
 class UUIDPrimaryKeyModel(models.Model):
@@ -21,6 +26,7 @@ class Container(AbstractEAVContainer):
 
 class Patient(UUIDPrimaryKeyModel):
     name = models.CharField(max_length=12)
+    label = models.CharField(max_length=12, default="", null=True)
     example = models.ForeignKey("tests.ExampleModel", null=True, blank=True, on_delete=models.PROTECT)
     container = models.ForeignKey("tests.Container", on_delete=models.CASCADE)
 
@@ -52,6 +58,40 @@ def __unicode__(self):
         return self.name
 
 
+class CustomQuerySet(EavQuerySet):
+
+    def name_search(self, value):
+        return self.filter(name=value)
+
+
+class CustomManager(EntityManager):
+
+    def name_search(self, value):
+        return self.filter(name=value)
+
+
+@register_eav()
+class ExampleModelWithCustomQuerySet(UUIDPrimaryKeyModel):
+    container = models.ForeignKey("tests.Container", on_delete=models.CASCADE)
+    name = models.CharField(max_length=12)
+
+    def __unicode__(self):
+        return self.name
+
+    objects = CustomQuerySet.as_manager()
+
+
+@register_eav()
+class ExampleModelWithCustomManager(UUIDPrimaryKeyModel):
+    container = models.ForeignKey("tests.Container", on_delete=models.CASCADE)
+    name = models.CharField(max_length=12)
+
+    def __unicode__(self):
+        return self.name
+
+    objects = CustomManager()
+
+
 @register_eav(container_field="custom_field")
 class ExampleModelCustomContainer(UUIDPrimaryKeyModel):
     custom_field = models.ForeignKey("tests.Container", on_delete=models.CASCADE)

tests/test_forms.py

@@ -3,7 +3,7 @@
 import pytest
 from django.contrib.admin.sites import AdminSite
 from django.core.handlers.base import BaseHandler
-from django.forms import ModelForm
+from django.forms import ModelForm, fields
 from django.test.client import RequestFactory
 
 import eav
@@ -46,6 +46,8 @@ def has_perm(self, perm):
 
 
 class PatientForm(ModelForm):
+    label = fields.CharField(required=False)
+
     class Meta:
         model = Patient
         fields = "__all__"
@@ -70,7 +72,7 @@ def test_fields(self, patient_with_short_attributes):
         admin = BaseEntityAdmin(Patient, self.site)
         admin.form = BaseDynamicEntityForm
         view = admin.change_view(request, str(patient_with_short_attributes.pk))
-        own_fields = 3
+        own_fields = 4
         adminform = view.context_data["adminform"]
         assert len(adminform.form.fields) == patient_with_short_attributes.container.attributes.count() + own_fields
 

tests/test_queries.py

@@ -6,7 +6,7 @@
 import eav
 from eav.models import Attribute, Value
 
-from .models import Encounter, Patient
+from .models import Encounter, Patient, ExampleModelWithCustomManager
 
 pytestmark = [pytest.mark.django_db]
 
@@ -69,6 +69,13 @@ def test_get_or_create_with_eav(self, container_with_short_attributes):
         assert Patient.objects.count() == 2
         assert Value.objects.count() == 2
 
+        Patient.objects.get_or_create(
+            defaults={"label": "test"}, fields__age=7, name="Bob", container_id=container_with_short_attributes.id
+        )
+        assert Patient.objects.count() == 3
+        assert Patient.objects.get(fields__age=7, name="Bob").label == "test"
+        assert Value.objects.count() == 3
+
     def test_get_with_eav(self, patient_with_extended_attributes):
         patient_with_extended_attributes.fields.age = 6
         patient_with_extended_attributes.save()
@@ -190,3 +197,13 @@ def order(*ordering, **kwargs):
             "down attributes to specific container."
         ):
             Patient.objects.all().order_by("fields__country")
+
+    def test_custom_queryset_search(self, example_model_with_custom_manager_factory):
+        example_model_with_custom_manager_factory(name="name1")
+        example_model_with_custom_manager_factory(name="name2")
+        assert ExampleModelWithCustomManager.objects.name_search("name1").count() == 1
+
+    def test_custom_manager_search(self, example_model_with_custom_manager_factory):
+        example_model_with_custom_manager_factory(name="name1")
+        example_model_with_custom_manager_factory(name="name2")
+        assert ExampleModelWithCustomManager.objects.name_search("name1").count() == 1

tests/test_registry.py

@@ -2,7 +2,7 @@
 import eav
 from eav.registry import EavConfig
 
-from .models import Encounter, ExampleModel, Patient
+from .models import Encounter, ExampleModel, Patient, ExampleModelWithCustomManager
 
 
 class TestRegistry:
@@ -74,3 +74,6 @@ def test_doesnt_register_nonmodel(self):
             @eav.decorators.register_eav()
             class Foo(object):
                 pass
+
+    def test_custom_manager(self):
+        assert hasattr(ExampleModelWithCustomManager.objects, "name_search")