django-eav2/eav/managers.py

"""
This module contains the custom manager used by entities registered with eav.
"""

from django.db import models

from eav.queryset import EavQuerySet


class EntityManager(models.Manager):
    """
    Our custom manager, overrides ``models.Manager``.
    """

    _queryset_class = EavQuerySet

    def create(self, **kwargs):
        """
        Parse eav attributes out of *kwargs*, then try to create and save
        the object, then assign and save it's eav attributes.
        """
        config_cls = getattr(self.model, "_eav_config_cls", None)

        if not config_cls or config_cls.manager_only:
            return super().create(**kwargs)

        prefix = f"{config_cls.eav_attr}__"
        new_kwargs = {}
        eav_kwargs = {}

        for key, value in kwargs.items():
            if key.startswith(prefix):
                eav_kwargs.update({key[len(prefix) :]: value})
            else:
                new_kwargs.update({key: value})

        obj = self.model(**new_kwargs)
        obj_eav = getattr(obj, config_cls.eav_attr)

        for key, value in eav_kwargs.items():
            setattr(obj_eav, key, value)

        obj.save()
        return obj

    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:
            if defaults:
                kwargs = {**kwargs, **defaults}
            return self.create(**kwargs), True