django-taggit TagAdder

django_taggit_tag_adder.py

from typing import List, Set, Tuple, Type
from taggit.models import Tag, TaggedItem
from django.db.models import Model, QuerySet
from django.contrib.contenttypes.models import ContentType


class TagAdder:
    """
    An object which adds specific tags to given Django Model objects in bulk.
    """
    _tags_exist: bool = False

    def _set_tag_ids(self, tags: QuerySet):
        """
        Args:
            tags (QuerySet[Tag]):
        """
        self._tags_exist: bool = True
        self.tag_ids: Tuple[int] = tuple(tags.values_list('pk', flat=True))

    def __init__(self, tag_names: Set[str], django_model: Type[Model]) -> None:
        """
        Args:
            tag_names (Set[str]):
                The names of the tags which should be added to given objects.
            django_model (Type[Model]):
                The type of Django Model whose objects will be assigned tags.
        """
        self.django_model = django_model
        if not hasattr(django_model, 'tags'):
            raise AttributeError(f'{django_model} has no attribute `tags`.')
        self.content_type_id: int = \
            ContentType.objects.get_for_model(model=django_model).pk
        self._tag_names: Set[str] = tag_names

        # TODO: what about partial matches? i.e. two of three tags exist?
        tags: QuerySet = Tag.objects.filter(name__in=tag_names)
        existing_tag_difference: Set[str] = \
            tag_names.difference(tag.name for tag in tags)
        if not existing_tag_difference:
            self._set_tag_ids(tags)

        def __call__(self, objects: Union[QuerySet, List]) -> List:
        """
        Set tags on objects; creates tags first if they do not already exist.

        Args:
            objects (Union[QuerySet, List]): Django objects which should be tagged.
        """
        first = objects[0]

        # TODO: is the `isinstance` type check really necessary? It adds some overhead
        if not isinstance(first, self.django_model):
            raise ValueError(
                f'Object {first} is not of type {self.django_model}. '
                f'Expected {self.django_model} but received {type(first)} instead.'
            )
        if not self._tags_exist:
            first.tags.set(*self._tag_names)
            self._set_tag_ids(first.tags.filter(name__in=self._tag_names))
            return self._set_tags(objects[1:])
        else:
            return self._set_tags(objects)

    def _set_tags(self, objects: Union[QuerySet, List]) -> List:
        """
        Set tags on objects via a ManyToMany `through` bulk_create.

        Args:
            objects (Union[QuerySet, List]):
        """
        tag_throughs: List[TaggedItem] = [
            self.django_model.tags.through(
                tag_id=tag_id,
                object_id=obj.id,
                content_type_id=self.content_type_id,
            )
            for obj in objects
            for tag_id in self.tag_ids
        ]

        return self.django_model.tags.through.objects.bulk_create(tag_throughs)